packages feed

futhark 0.26.4 → 0.27.1

raw patch · 138 files changed

+21186/−8546 lines, 138 filesdep +OneTupledep ~futhark-datadep ~regex-tdfadep ~srcloc

Dependencies added: OneTuple

Dependency ranges changed: futhark-data, regex-tdfa, srcloc

Files

CHANGELOG.md view
@@ -5,12 +5,80 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/) and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.html). +## [0.27.1]++## Added++* Rewrote partition to be a single kernel.++* The type checker has been rewritten, with contributions from Jacob Aleksandar+  Siegumfeldt, Laust Kjæp Dengsøe, and Robert Schenck.++* Flattening has been rewritten, with the majority of the work by Amirreza+  Hashemi. The main consequence is that nonuniform nested parallelism is now+  supported, although it is not yet necessarily efficient in all cases. More+  details: https://futhark-lang.org/blog/2026-07-31-full-flattening.html++* Futhark now supports recursive functions, with various restrictions.++* The `incremental_flattening` attributes are now just named `flattening`,+  although the old name continues to work.++* A simplification rule for branches that return common results (#2526).++## Changed++* Local functions are no longer let-generalised (i.e., made polymorphic).+  Explicitly polymorphic local functions are still supported.++## Fixed++* Filter now uses the predicate once per element instead of twice.++* Fixed bug in simiplification engine, various SOACs needed to specify+  the depth of lambdas, scans, and reduces.++* In-place updates with a slice that covers the entire array, but reorders its+  elements (such as a reversal), were simplified into a copy, discarding the+  reordering. Among other things this produced wrong gradients for `reverse`+  (#2522).++* A case where complex sizes referring to explicit parameters were mishandled by+  monomorphisation (#2230).++* An issue where `#[scratch]` would apply to subexpressions in undesirable ways,+  changing the type of the expression at the IR level.++* Simplified fusibility check by removing redundant accumulator overlap check+  and fixed fusibility check by giving the correct number of elements.++* A compiler crash due to missing double buffering inside sequential code+  migrated to GPU. (#2513)++* Atomics now use device-wide memory scope when compiling for OpenCL C 2.0, or+  OpenCL C 3.0 with the required atomic features. OpenCL C 2.0 is selected+  automatically for Rusticl on Asahi; elsewhere it must be selected explicitly.+  This fixes incorrect results from cross-workgroup operations+  ([#734](https://github.com/diku-dk/futhark/issues/734)).++* An exotic case in reverse-mode differentiation of accumulators, which in+  practice would only occur in cases of unrolled `scatter`s.+ ## [0.26.4]  ### Added  * `futhark repl` has a new command: `:string`. +* `futhark benchcmp` has new options `--sort-by` and `--order` for+  controlling the order in which program groups are printed.+  `--sort-by=significant` sorts by number of statistically significant+  regressions, `--sort-by=geomean-significant` and+  `--sort-by=geomean-all` sort by geometric mean speedup over+  significant or all datasets respectively.  `--order=worst-first`+  (default) surfaces regressions at the top; `--order=best-first`+  surfaces improvements.  The default remains unsorted (alphabetical).+ * The `hip` backend previously simulated `f16` operations with `f32`, but now it   uses the hardware support for `f16`, similarly to the CUDA backend.   Implemented by Jérôme Wagner. (#2470)@@ -543,7 +611,7 @@  * `futhark profile` now supports multiple JSON files. -* `futhark fmt`, by William Due and Therese Lyngby.+* `futhark fmt`, by Lilje Due and Therese Lyngby.  * Lambdas can now be passed as the last argument to a function application. @@ -1227,7 +1295,7 @@ * Datasets used in `futhark test` and `futhark bench` can now be named   (#1859). -* New command `futhark benchcmp` by William Due.+* New command `futhark benchcmp` by Lilje Due.  ### Changed
docs/conf.py view
@@ -506,6 +506,13 @@         1,     ),     (+        "man/futhark-benchcmp",+        "futhark-benchcmp",+        "compare Futhark benchmark results",+        [],+        1,+    ),+    (         "man/futhark-doc",         "futhark-doc",         "generate documentation for Futhark code",
docs/error-index.rst view
@@ -184,41 +184,10 @@   def f n =     g (iota n, iota n)) -.. _consuming-parameter:--"Consuming parameter passed non-unique argument"-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~--Caused by programs like the following:--.. code-block:: futhark--  def update (xs: *[]i32) = xs with [0] = 0--  def f (ys: []i32) = update ys--The update ``function`` *consumes* its ``xs`` argument to perform an-:ref:`in-place update <in-place-updates>`, as denoted by the asterisk-before the type.  However, the ``f`` function tries to pass an array-that it is not allowed to consume (no asterisk before the type).--One solution is to change the type of ``f`` so that it also consumes-its input, which allows it to pass it on to ``update``:--.. code-block:: futhark--  def f (ys: *[]i32) = update ys--Another solution to ``copy`` the array that we pass to ``update``:--.. code-block:: futhark--  def f (ys: []i32) = update (copy ys)- .. _consuming-argument: -"Non-consuming higher-order parameter passed consuming argument."-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~+"Argument of functional type ... contains consumption"+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~  This error occurs when we have a higher-order function that expects a function that does *not* consume its arguments, and we pass it one@@ -285,26 +254,6 @@ Therefore, the type checker invents an :term:`unknown size` variable, say ``l``, and assigns ``a`` the type ``[l]i32``. -.. _size-expression-consume:--"Size expression with consumption is replaced by unknown size."-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~--To illustrate this error, consider the following program--.. code-block:: futhark--   def consume (xs: *[]i64): i64 = xs[0]--   def main (xs: *[]i64) =-     let a = iota (consume xs)-     in ...--Intuitively, the type of ``a`` should be ``[consume ys]i32``, but this-puts a consumption of the array ``ys`` into a size expression, which-is invalid.  Therefore, the type checker invents an :term:`unknown-size` variable, say ``l``, and assigns ``a`` the type ``[l]i32``.- .. _inaccessible-size:  "Parameter *x* refers to size *y* which will not be accessible to the caller@@ -367,10 +316,42 @@ Here the type rules force ``A`` to have size ``x``, leading to a problematic type.  It can be fixed using the techniques above. +.. _consuming-loop-param-aliases:++"Return value for consuming loop parameter *x* aliases *y*"+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~++This occurs for expressions like the following::++    loop (xs: []i32, ys: *[]i32) = (replicate n 0, replicate n 0)+    for i < 10 do+      (xs, xs)+++This is not allowed, as creates aliasing between a consumeable parameter+(``ys``) and non-consumable parameter (``xs``) in the next iteration of the+loop, during which consumption ``ys`` would also affect ``xs``. You can solve+this by copying one of the return values of the loop.++.. _loop-parameter-aliases-other:++"Return value for loop parameter *x* aliases other consumed loop parameter"+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~++This occurs for expressions like the following::++    loop (xs: *[]i32, ys: *[]i32) = (replicate n 0, replicate n 0)+    for i < 10 do+      (xs, xs)++This is not allowed for the same reason that we are not allowed to consume an+array multiple times. You can solve this by copying one of the return values of+the loop.+ .. _aliases-previously-returned:  "Return value for consuming loop parameter *x* aliases previously returned value"----------------------------------------------------------------------------------+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~  This error occurs when you have a loop with multiple loop parameters, at least one of which is consuming, and the values returned by the@@ -388,7 +369,28 @@     -- is consumed.     in (arr[i+1], arr) +.. _contains-consumption: +"Let-bound expression of higher-order type *t* contains consumption"+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~++This occurs when ``let``-binding an expression that contains consumption and+returns a function. The most common case is partial application of a consuming+function::++  def update (xs: *[]i32) (i: i32) (y: i32) =+    xs with [i] = y++  def main (xs: *[]i32) =+    let f = update xs+    in (f 0 0, f 0 0)++The simplest solution is to remove the consumption by doing a ``copy``.++The reason for this restriction is rooted in efficiency concerns.+Defunctionalisation causes the two applications of ``f`` to both consume ``xs``,+which is a violation of uniqueness properties.+ Size errors ----------- @@ -720,6 +722,26 @@ Other errors ------------ +.. _scope-violation:++"Scope violation"+~~~~~~~~~~~~~~~~~++These occurs when the type checker infers that the type (or size) of a variable+``x`` is forced to be expressed using parameters or variables not in scope when+``x`` is bound. This only occurs when mixing explicit parameters with inference.+Contrived example::++  def f x =+    let g 'b (y: b) = if true then y else x+    in g++The ``if`` forces ``y`` and ``x`` to be the same type, but ``y`` has type ``b``,+which is a type parameter bound in ``g``, and not in scope where ``x`` is bound.++These errors usually imply some form of misdesign, and can be resolved by+manually inserting type annotations until the conceptual mistake becomes clear.+ .. _literal-out-of-bounds:  "Literal out of bounds"@@ -878,6 +900,22 @@ .. code-block:: futhark    def f (r : {x:i32}) = r with x = 0++.. _occurs-check:++"Occurs check"+~~~~~~~~~~~~~~++Occurs check errors are reported whenever the type checker infers a type or size+that is circular, meaning it would be infinitely large. Essentially, whenever+the type checker infers that some ``x`` must be equal to some (compound)+construct ``y``, it checks whether ``x`` is present inside ``y``. This is called+the occurs check.++Since it is such a general mechanism, there is no rule of thumb for how to avoid+or fix errors that manifest as an occurs check. Since they are always the result+of a misdesign, it can be useful to add explicit type annotations until the root+cause is revealed.  Entry points ------------
docs/glossary.rst view
@@ -136,10 +136,8 @@     GPU backend -     A :term:`compiler backend` that ultimately produces GPU code.-     The backends ``opencl`` and ``gpu`` are GPU backends.  These have-     more restrictions than some other backends, particularly with-     respect to :term:`irregular nested data parallelism`.+     A :term:`compiler backend` that ultimately produces GPU code. The backends+     ``opencl`` and ``gpu`` are GPU backends.     Higher-ranked type @@ -165,53 +163,12 @@      example, ``(x,y)`` is a pattern that will match any tuple. See      also :term:`refutable pattern`. -   Irregular--     Something that is not regular.  Usually used as shorthand for-     :term:`irregular nested data parallelism` or :term:`irregular-     array`.-    Irregular array       An array where the elements do not have the same size.  For      example, ``[[1], [2,3]`` is irregular.  These are not supported      in Futhark. -   Irregular nested data parallelism--     An instance of :term:`nested data parallelism`, where the-     :term:`parallel width` of inner parallelism is :term:`variant` to-     the outer parallelism.  For example, the following expression-     exhibits irregular nested data parallelism::--       map (\n -> reduce (+) 0 (iota n)) ns--     Because the width of the inner ``reduce`` is ``n``, and every-     iteration of the outer ``map`` has a (potentially) different-     ``n``.  The Futhark :term:`GPU backends<GPU backend>` *currently*-     do not support irregular nested data parallelism well, and will-     usually sequentialise the irregular loops.  In cases that require-     an :term:`irregular memory allocation`, the compiler may entirely-     fail to generate code.--   Irregular memory allocation--     A situation that occurs when the generated code has to allocate-     memory inside of an instance of :term:`nested data parallelism`,-     where the amount to allocate is variant to the outer parallel-     levels.  As a contrived example (that the actual compiler would-     just optimise away), consider::--       map (\n -> let A = iota n-                  in A[10])-           ns--     To construct the array ``A`` in memory, we require ``8n`` bytes,-     but ``n`` is not known until we start executing the body of the-     ``map``.  While such simple cases are handled, more complicated-     ones that involve nested sequential loops are not supported by-     the :term:`GPU backends<GPU backend>`.-    Parametric module       A function from :term:`modules<module>` to modules.  The most@@ -239,6 +196,12 @@      themselves.  These have various restrictions on their use in      order to support :term:`defunctionalisation`.  See :ref:`hofs`. +   Map nest++     The slightly inaccurate term for a collection of nested parallel+     operations. We use this term even when not all of the parallel dimensions+     are actually ``map``.+    Module       A mapping from names to definitions of types, values, or nested@@ -287,6 +250,43 @@      is used inside of another parallel construct.  For example, a      ``reduce`` might be used inside a function passed to ``map``. +   Nonuniform++     Something that is not :term:`uniform`. Usually used as shorthand for+     :term:`nonuniform nested data parallelism`.++   Nonuniform nested data parallelism++     An instance of :term:`nested data parallelism`, where the :term:`parallel+     width` of inner parallelism is :term:`variant` to the outer parallelism.+     For example, the following expression exhibits nonuniform nested data+     parallelism::++       map (\n -> reduce (+) 0 (iota n)) ns++     Because the width of the inner ``reduce`` is ``n``, and every iteration of+     the outer ``map`` has a (potentially) different ``n``. This is+     substantially more difficult to handle than :term:`uniform nested data+     parallelism` and comes with a nontrivial runtime cost.++   Nonuniform memory allocation++     A situation that occurs when the generated code has to allocate+     memory inside of an instance of :term:`nested data parallelism`,+     where the amount to allocate is variant to the outer parallel+     levels.  As a contrived example (that the actual compiler would+     just optimise away), consider::++       map (\n -> let A = iota n+                  in A[10])+           ns++     To construct the array ``A`` in memory, we require ``8n`` bytes,+     but ``n`` is not known until we start executing the body of the+     ``map``.  While such simple cases are handled, more complicated+     ones that involve nested sequential loops are not supported by+     the :term:`GPU backends<GPU backend>`.+    Parallel width       A somewhat informal term used to describe the size of an array on@@ -315,11 +315,6 @@      expressions or in function parameters. See also      :term:`irrefutable pattern`. -   Regular nested data parallelism--     An instance of :term:`nested data parallelism` that is not-     :term:`irregular`.  Fully supports by any :term:`GPU backend`.-    Size       The symbolic size of an array dimension or :term:`abstract type`.@@ -421,6 +416,20 @@      rest of the definition.  Do not confuse them with :term:`size      parameters <size parameter>`. +   Uniform++     Whether something is invariant :term:`invariant` to the enclosing+     :term:`map nest`. The precise meaning depends on the thing is question, for+     example, we speak of :term:`uniform nested data parallelism` if the size+     (and control flow leading to) a parallel construct is invariant to the+     enclosing map nest, without considering the concrete values.++   Uniform nested data parallelism++     An instance of :term:`nested data parallelism` that is not+     :term:`nonuniform`. This is much more efficient than :term:`nonuniform+     nested data parallelism`.+    Uniqueness types       A somewhat misleading term that describes Futhark's system of@@ -446,9 +455,8 @@     Variant -     When some value ``v`` computed inside a loop takes a different-     value for each iteration inside the loop, we say that ``v`` is-     *variant* to the loop (and otherwise :term:`invariant`).  Often-     used to talk about :term:`irregularity <irregular>`.  When-     something is nested inside multiple loops, it may be variant to-     just one of them.+     When some value ``v`` computed inside a loop takes a different value for+     each iteration inside the loop, we say that ``v`` is *variant* to the loop+     (and otherwise :term:`invariant`). Often used to talk about+     :term:`uniformity <uniform>`. When something is nested inside multiple+     loops, it may be variant to just one of them.
docs/index.rst view
@@ -47,6 +47,7 @@     man/futhark-autotune.rst    man/futhark-bench.rst+   man/futhark-benchcmp.rst    man/futhark-c.rst    man/futhark-cuda.rst    man/futhark-dataset.rst
docs/language-reference.rst view
@@ -291,11 +291,11 @@    def name params...: rettype = body -Hindley-Milner-style type inference is supported.  A parameter may be-given a type with the notation ``(name: type)``.  Functions may not be-recursive.  The sizes of the arguments can be constrained - see `Size-Types`_.  A function can be *polymorphic* by using type parameters, in-the same way as for `Type Abbreviations`_::+Hindley-Milner-style type inference is supported. A parameter may be given a+type with the notation ``(name: type)``. Functions may be recursive, subject to+various restrictions - see :ref:`recursive-functions`. The sizes of the arguments+can be constrained - see `Size Types`_. A function can be *polymorphic* by using+type parameters, in the same way as for `Type Abbreviations`_::    def reverse [n] 't (xs: [n]t): [n]t = xs[::-1] @@ -1006,9 +1006,11 @@ ...............................  Bind ``f`` to a function with the given parameters and definition-(``e``) and evaluate ``body``.  The function will be treated as-aliasing any free variables in ``e``.  The function is not in scope of-itself, and hence cannot be recursive.+(``e``) and evaluate ``body``. The function will be treated as+aliasing any free variables in ``e``. The function is not in scope of+itself, and hence cannot be recursive. While the function can be made+polymorphic by putting in explicit size parameters, it is not+automatically generalised the way top level functions are.  ``loop pat = initial for x in a do loopbody`` .............................................@@ -1105,6 +1107,9 @@  * A ``loop`` parameter cannot be a function. +* There are some restrictions on recursive higher-order functions, see+  :ref:`recursive-functions`.+ Further, *type parameters* are divided into *non-lifted* (bound with an apostrophe, e.g. ``'t``), *size-lifted* (``'~t``), and *fully lifted* (``'^t``).  Only fully lifted type parameters may be@@ -1135,6 +1140,31 @@ let-generalisation *unless* they are syntactically functions, meaning they have at least one named parameter. +.. _recursive-functions:++Recursive functions+-------------------++Functions may be recursive, subject to the following constraints.++1. Mutual recursion is not supported - a function is only in scope of itself.++2. A recursive function definition may not have a higher-order return type. In+   some cases you can work around this restriction by adding more function+   parameters.++3. A recursive call of a higher-order function must be provided arguments for+   all of its higher-order parameters, and they must be syntactically identical+   to the corresponding parameter in the definition.++Restriction 2 is to make restriction 3 feasible to check, and restriction 3+exists to ensure defunctionalisation is possible.++Recursion is monomorphic, meaning a recursive call uses the same instantiation+of type parameters as the function definition. However, *size-polymorphic*+recursion is allowed, but requires explicit size parameters in the function+definition.+ .. _size-types:  Size Types@@ -1772,28 +1802,42 @@ This is used to work around optimisation deficiencies (or bugs), although it should hopefully rarely be necessary. -``incremental_flattening(no_outer)``-....................................+``flattening(sequentialise_nonuniform)``+........................................ +Within the attributed SOAC (which should not itself be nested), any nonuniform+nested parallelism is sequentialised. This gives up the inner parallelism, but+avoids the substantial bookkeeping that flattening irregular arrays entails.+Note that in some cases this can lead to un-compileable code, if the+nonuniformity is too severe. Uniform nested parallelism is unaffected. The+attribute is ignored where the irregularity cannot be confined to a single+thread, in which case such values are flattened as usual.++Historical note: this is intended to reproduce the behaviour of Futhark prior to+supporting irregular flattening.++``flattening(no_outer)``+........................+ When using incremental flattening, do not generate the "only outer parallelism" version for the attributed SOACs. -``incremental_flattening(no_intra)``-....................................+``flattening(no_intra)``+........................  When using incremental flattening, do not generate the "intra-block parallelism" version for the attributed SOACs. -``incremental_flattening(only_intra)``-......................................+``flattening(only_intra)``+..........................  When using incremental flattening, *only* generate the "intra-block parallelism" version of the attributed SOACs.  **Beware**: the resulting program will fail to run if the inner parallelism does not fit on the device. -``incremental_flattening(only_inner)``-......................................+``flattening(only_inner)``+..........................  When using incremental flattening, do not generate multiple versions for this SOAC, but do exploit inner parallelism (which may give rise
docs/man/futhark-bench.rst view
@@ -200,4 +200,4 @@ SEE ALSO ======== -:ref:`futhark-c(1)`, :ref:`futhark-test(1)`+:ref:`futhark-c(1)`, :ref:`futhark-test(1)`, :ref:`futhark-benchcmp(1)`
+ docs/man/futhark-benchcmp.rst view
@@ -0,0 +1,93 @@+.. role:: ref(emphasis)++.. _futhark-benchcmp(1):++================+futhark-benchcmp+================++SYNOPSIS+========++futhark benchcmp [options...] FILE_A FILE_B++DESCRIPTION+===========++Compare two JSON files produced by the ``--json`` option of+:ref:`futhark-bench(1)` and print a human-readable summary of the+speedup of ``FILE_B`` relative to ``FILE_A``.  A speedup greater than+1x means ``FILE_B`` is faster; a speedup less than 1x means ``FILE_B``+is slower (a regression).++Results are grouped by program and entry point.  Within each group the+datasets are listed in alphabetical order.  A speedup is highlighted in+green when it is statistically significant and faster, and in red when+it is statistically significant and slower.  Significance is determined+by comparing the difference in means against the sum of half-standard-+deviations of the two samples.++``FILE_A`` is conventionally the *baseline* and ``FILE_B`` the *new*+result.  The typical workflow is::++  futhark bench --backend=cuda --json baseline.json prog.fut+  # ... make changes ...+  futhark bench --backend=cuda --json new.json prog.fut+  futhark benchcmp baseline.json new.json++OPTIONS+=======++--sort-by=METRIC++  Sort program groups by the given metric.  The default is unsorted+  (alphabetical by program name, matching the order of ``futhark+  bench`` output).++  ``significant``+    Sort by the number of datasets in the group that have a+    statistically significant regression (speedup < 0.99).  Groups with+    the most regressions appear first.++  ``geomean-significant``+    Sort by the geometric mean of speedups restricted to statistically+    significant datasets.  Groups whose significant datasets are slowest+    on average appear first.  Groups with no significant results are+    treated as 1.0x (no change) for sorting purposes.++  ``geomean-all``+    Sort by the geometric mean of speedups across *all* datasets in the+    group, regardless of significance.  Groups that are slowest on+    average appear first.++--order=ORDER++  Control the sort direction.  Only meaningful when ``--sort-by`` is+  also given.++  ``worst-first`` (default)+    Surface the most regressed programs at the top.  For+    ``significant`` this means the highest count first; for the geomean+    metrics it means the lowest ratio first (since a ratio below 1+    indicates a slowdown).++  ``best-first``+    Surface the most improved programs at the top.++EXAMPLES+========++Compare two benchmark runs and show the worst regressions first by+number of significant datasets::++  futhark benchcmp --sort-by=significant baseline.json new.json++Show the same comparison ordered by the geometric mean over all+datasets, with the most improved programs first::++  futhark benchcmp --sort-by=geomean-all --order=best-first baseline.json new.json++SEE ALSO+========++:ref:`futhark-bench(1)`
docs/man/futhark.rst view
@@ -24,12 +24,13 @@ COMMANDS ======== -futhark benchcmp FILE_A FILE_B-------------------------------+futhark benchcmp [options...] FILE_A FILE_B+-------------------------------------------  Compare two JSON files produced by the ``--json`` option of-:ref:`futhark-bench(1)`.  The results show speedup of the latter file-compared to the former.+:ref:`futhark-bench(1)`.  ``FILE_A`` is the baseline and ``FILE_B`` is+the new result; speedups greater than 1x mean ``FILE_B`` is faster.+See :ref:`futhark-benchcmp(1)` for the full option reference.  futhark check [-w] [-Werror] PROGRAM ------------------------------------@@ -116,4 +117,4 @@ SEE ALSO ======== -:ref:`futhark-opencl(1)`, :ref:`futhark-c(1)`, :ref:`futhark-py(1)`, :ref:`futhark-pyopencl(1)`, :ref:`futhark-wasm(1)`, :ref:`futhark-wasm-multicore(1)`, :ref:`futhark-ispc(1)`, :ref:`futhark-dataset(1)`, :ref:`futhark-doc(1)`, :ref:`futhark-test(1)`, :ref:`futhark-bench(1)`, :ref:`futhark-run(1)`, :ref:`futhark-repl(1)`, :ref:`futhark-literate(1)`+:ref:`futhark-opencl(1)`, :ref:`futhark-c(1)`, :ref:`futhark-py(1)`, :ref:`futhark-pyopencl(1)`, :ref:`futhark-wasm(1)`, :ref:`futhark-wasm-multicore(1)`, :ref:`futhark-ispc(1)`, :ref:`futhark-dataset(1)`, :ref:`futhark-doc(1)`, :ref:`futhark-test(1)`, :ref:`futhark-bench(1)`, :ref:`futhark-benchcmp(1)`, :ref:`futhark-run(1)`, :ref:`futhark-repl(1)`, :ref:`futhark-literate(1)`
docs/versus-other-languages.rst view
@@ -84,12 +84,12 @@  Lambda terms are written as ``\x -> x + 2``, as in Haskell. -A Futhark program is read top-down, and all functions must be declared-in the order they are used, like Standard ML.  Unlike just-about all functional languages, recursive functions are *not*-supported.  Most of the time, you will use bulk array operations-instead, but there is also a dedicated ``loop`` language construct,-which is essentially syntactic sugar for tail recursive functions.+A Futhark program is read top-down, and all functions must be declared in the+order they are used, like Standard ML. Although Futhark does support recursive+functions with some restrictions (see :ref:`recursive-functions`), their use is+discouraged. Most of the time, you will use bulk array operations instead, but+there is also a dedicated ``loop`` language construct, which is essentially+syntactic sugar for tail recursive functions.  Types -----
futhark.cabal view
@@ -1,6 +1,6 @@ cabal-version: 3.0 name:           futhark-version:        0.26.4+version:        0.27.1 synopsis:       An optimising compiler for a functional, array-oriented language.  description:    Futhark is a small programming language designed to be compiled to@@ -368,19 +368,24 @@       Futhark.Pass.ExplicitAllocations.MC       Futhark.Pass.ExplicitAllocations.SegOp       Futhark.Pass.ExplicitAllocations.Seq-      Futhark.Pass.ExtractKernels-      Futhark.Pass.ExtractKernels.BlockedKernel-      Futhark.Pass.ExtractKernels.DistributeNests-      Futhark.Pass.ExtractKernels.Distribution-      Futhark.Pass.ExtractKernels.ISRWIM-      Futhark.Pass.ExtractKernels.Interchange-      Futhark.Pass.ExtractKernels.Intrablock-      Futhark.Pass.ExtractKernels.StreamKernel-      Futhark.Pass.ExtractKernels.ToGPU       Futhark.Pass.ExtractMulticore       Futhark.Pass.FirstOrderTransform+      Futhark.Pass.Flatten+      Futhark.Pass.Flatten.BasicOp+      Futhark.Pass.Flatten.Builtins+      Futhark.Pass.Flatten.PreProcess+      Futhark.Pass.Flatten.Incremental+      Futhark.Pass.Flatten.Intrablock+      Futhark.Pass.Flatten.Distribute+      Futhark.Pass.Flatten.General+      Futhark.Pass.Flatten.Loop+      Futhark.Pass.Flatten.Match+      Futhark.Pass.Flatten.Monad+      Futhark.Pass.Flatten.SOAC+      Futhark.Pass.Flatten.WithAcc       Futhark.Pass.LiftAllocations       Futhark.Pass.LowerAllocations+      Futhark.Pass.NoGrid       Futhark.Pass.Simplify       Futhark.Passes       Futhark.Pipeline@@ -400,8 +405,10 @@       Futhark.Tools       Futhark.Transform.CopyPropagate       Futhark.Transform.FirstOrderTransform+      Futhark.Transform.ISRWIM       Futhark.Transform.Rename       Futhark.Transform.Substitute+      Futhark.Transform.ToGPU       Futhark.Util       Futhark.Util.CMath       Futhark.Util.IntegralExp@@ -435,6 +442,10 @@       Language.Futhark.Tuple       Language.Futhark.TypeChecker       Language.Futhark.TypeChecker.Consumption+      Language.Futhark.TypeChecker.Constraints+      Language.Futhark.TypeChecker.TySolve+      Language.Futhark.TypeChecker.UnionFind+      Language.Futhark.TypeChecker.Error       Language.Futhark.TypeChecker.Names       Language.Futhark.TypeChecker.Match       Language.Futhark.TypeChecker.Modules@@ -443,6 +454,8 @@       Language.Futhark.TypeChecker.Terms.Loop       Language.Futhark.TypeChecker.Terms.Monad       Language.Futhark.TypeChecker.Terms.Pat+      Language.Futhark.TypeChecker.Terms.Scope+      Language.Futhark.TypeChecker.Terms.Unsized       Language.Futhark.TypeChecker.Types       Language.Futhark.TypeChecker.Unify       Language.Futhark.Warnings@@ -479,7 +492,7 @@     , file-embed >=0.0.14.0     , filepath >=1.4.1.1     , free >=5.1.10-    , futhark-data >= 1.1.3.0+    , futhark-data >= 1.1.4.0     , futhark-server >= 1.4.1.0     , futhark-manifest == 1.9.0.0     , githash >=0.1.6.1@@ -491,6 +504,7 @@     , lsp-types >= 2.4.0.0     , mainland-pretty >=0.7.1     , cmark-gfm >=0.2.1+    , OneTuple     , megaparsec >=9.0.0     , mtl >=2.2.1     , neat-interpolation >=0.3@@ -525,6 +539,8 @@   hs-source-dirs: src-testing   visibility: private   exposed-modules:+      Generated.AllFutBenchmarks+      Generated.AllFutBenchmarks.Accelerate.Nbody.Nbodybh       Futhark.AD.DerivativesTests       Futhark.Analysis.AlgSimplifyTests       Futhark.Analysis.DataDependenciesTests@@ -559,6 +575,8 @@       Language.Futhark.SemanticTests       Language.Futhark.SyntaxTests       Language.Futhark.TypeChecker.TypesTests+      Language.Futhark.TypeChecker.TySolveBenchmarks+      Language.Futhark.TypeChecker.TySolveTests       Language.Futhark.TypeChecker.ConsumptionTests       Language.Futhark.TypeChecker.ModulesTests       Language.Futhark.TypeCheckerTests@@ -575,13 +593,14 @@     , lsp     , lsp-test     , megaparsec+    , srcloc >=0.4     , neat-interpolation     , process-    , srcloc     , tasty     , tasty-hunit     , tasty-quickcheck     , text+    , regex-tdfa ^>= 1.3.2  test-suite unit   import: common
prelude/soacs.fut view
@@ -241,8 +241,7 @@ def filter [n] 'a (p: a -> bool) (as: [n]a) : *[]a =   let flags = map p as   let offsets = scan (+) 0 (map intrinsics.btoi_bool_i64 flags)-  let flags' = map p as-  let is = map2 (\f o -> if f then o - 1 else -1) flags' offsets+  let is = map2 (\f o -> if f then o - 1 else -1) flags offsets   -- This following is carefully written such that the two scatters will be   -- fused horisontally, which allows the entire thing to become a single   -- kernel.@@ -259,13 +258,15 @@ -- **Work:** *O(n ✕ W(p))* -- -- **Span:** *O(log(n) ✕ W(p))*-def partition [n] 'a (p: a -> bool) (as: [n]a) : ?[k].([k]a, [n - k]a) =-  let offset =-    reduce_comm (+) 0 (map (\x -> intrinsics.btoi_bool_i64 (p x)) as)+def partition [n] 'a+              (p: a -> bool)+              (as: [n]a) : ?[k].([k]a, [n - k]a) =+  let to_index0 f (o0, _o1) = if f then o0 - 1 else -1+  let to_index1 f (_o0, o1) = if f then n - o1 else -1   let add2 (a0, b0) (a1, b1) = (a0 + a1, b0 + b1)-  let to_index f (o0, o1) = if f then o0 - 1i64 else offset + o1 - 1   let t_flags = map p as-  let f_flags = map (\x -> !x) t_flags+  let rev_as = as[::-1]+  let f_flags = map (\x -> !x) (map p rev_as)   let flags =     map2 (\x y ->             ( intrinsics.btoi_bool_i64 x@@ -274,9 +275,15 @@          t_flags          f_flags   let offsets = scan add2 (0, 0) flags-  let idxs = map2 to_index (map p as) offsets-  let res = scatter (#[scratch] [as][0]) idxs as-  in (res[0:offset], res[offset:n])+  let idxs0 = map2 to_index0 t_flags offsets+  let idxs1 = map2 to_index1 f_flags offsets+  let is = intrinsics.concat idxs0 idxs1+  let res = scatter (#[scratch] [as][0]) is (intrinsics.concat as rev_as)+  let count =+    scatter [0]+            (map (\j -> if j == n - 1 then 0 else -1) (0..1..<n))+            (map (.0) offsets)+  in (res[0:count[0]], res[count[0]:n])  -- | Split an array by two predicates, producing three arrays. --@@ -314,3 +321,40 @@      , res[offset0:offset0 + offset1] :> [offset1]a      , res[offset0 + offset1:n] :> [n - offset0 - offset1]a      )++-- | Perform a flattened map of a function `f`, that produces a nonuniform and+-- uniform-sized result, over the array `as`. Returns the concatenation of the+-- nonuniform results and an array of the uniforms, alongside metadata allowing+-- the interpretation of the concatenated segments as an irregular array. In+-- order:+--+-- * The *shape array*, giving the size of each segment. This array sums to `m`.+--+-- * The *flag array*, indicating for each element when a new segment begins.+--+-- * The *offset array*, indicating for each segment where its values begin in+--   the data array.+--+-- * The *data array*, comprising the concatenated results of `f`.+--+-- * An array of uniform results, which has no special name.+def flatmap [n] 'a 'b 'c+            (f: a -> ?[k].([k]b, c))+            (as: [n]a) : ?[m].( [n]i64+                              , [m]bool+                              , [n]i64+                              , [m]b+                              , [n]c+                              ) =+  intrinsics.flatmap f as++-- | Like `flatmap`, but without the value result.+def flatmap' [n] 'a 'b+             (f: a -> ?[k].[k]b)+             (as: [n]a) : ?[m].( [n]i64+                               , [m]bool+                               , [n]i64+                               , [m]b+                               ) =+  let (S, F, O, D, _) = flatmap (\x -> (f x, ())) as+  in (S, F, O, D)
rts/c/atomics32.h view
@@ -28,6 +28,10 @@ SCALAR_FUN_ATTR int32_t atomic_xchg_i32_global(volatile __global int32_t *p, int32_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicExch((int32_t*)p, x);+#elif defined(FUTHARK_OPENCL_DEVICE_ATOMICS)+  return atomic_exchange_explicit((volatile atomic_int*)p, x,+                                  memory_order_acq_rel,+                                  memory_scope_device); #else   return atomic_xor(p, x); #endif@@ -45,6 +49,11 @@                                                   int32_t cmp, int32_t val) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicCAS((int32_t*)p, cmp, val);+#elif defined(FUTHARK_OPENCL_DEVICE_ATOMICS)+  atomic_compare_exchange_strong_explicit(+    (volatile atomic_int*)p, &cmp, val,+    memory_order_acq_rel, memory_order_acquire, memory_scope_device);+  return cmp; #else   return atomic_cmpxchg(p, cmp, val); #endif@@ -62,6 +71,10 @@ SCALAR_FUN_ATTR int32_t atomic_add_i32_global(volatile __global int32_t *p, int32_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicAdd((int32_t*)p, x);+#elif defined(FUTHARK_OPENCL_DEVICE_ATOMICS)+  return atomic_fetch_add_explicit((volatile atomic_int*)p, x,+                                   memory_order_acq_rel,+                                   memory_scope_device); #else   return atomic_add(p, x); #endif@@ -123,6 +136,10 @@ SCALAR_FUN_ATTR int32_t atomic_smax_i32_global(volatile __global int32_t *p, int32_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicMax((int32_t*)p, x);+#elif defined(FUTHARK_OPENCL_DEVICE_ATOMICS)+  return atomic_fetch_max_explicit((volatile atomic_int*)p, x,+                                  memory_order_acq_rel,+                                  memory_scope_device); #else   return atomic_max(p, x); #endif@@ -139,6 +156,10 @@ SCALAR_FUN_ATTR int32_t atomic_smin_i32_global(volatile __global int32_t *p, int32_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicMin((int32_t*)p, x);+#elif defined(FUTHARK_OPENCL_DEVICE_ATOMICS)+  return atomic_fetch_min_explicit((volatile atomic_int*)p, x,+                                  memory_order_acq_rel,+                                  memory_scope_device); #else   return atomic_min(p, x); #endif@@ -155,6 +176,10 @@ SCALAR_FUN_ATTR uint32_t atomic_umax_i32_global(volatile __global uint32_t *p, uint32_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicMax((uint32_t*)p, x);+#elif defined(FUTHARK_OPENCL_DEVICE_ATOMICS)+  return atomic_fetch_max_explicit((volatile atomic_uint*)p, x,+                                  memory_order_acq_rel,+                                  memory_scope_device); #else   return atomic_max(p, x); #endif@@ -171,6 +196,10 @@ SCALAR_FUN_ATTR uint32_t atomic_umin_i32_global(volatile __global uint32_t *p, uint32_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicMin((uint32_t*)p, x);+#elif defined(FUTHARK_OPENCL_DEVICE_ATOMICS)+  return atomic_fetch_min_explicit((volatile atomic_uint*)p, x,+                                  memory_order_acq_rel,+                                  memory_scope_device); #else   return atomic_min(p, x); #endif@@ -187,6 +216,10 @@ SCALAR_FUN_ATTR int32_t atomic_and_i32_global(volatile __global int32_t *p, int32_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicAnd((int32_t*)p, x);+#elif defined(FUTHARK_OPENCL_DEVICE_ATOMICS)+  return atomic_fetch_and_explicit((volatile atomic_int*)p, x,+                                  memory_order_acq_rel,+                                  memory_scope_device); #else   return atomic_and(p, x); #endif@@ -203,6 +236,10 @@ SCALAR_FUN_ATTR int32_t atomic_or_i32_global(volatile __global int32_t *p, int32_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicOr((int32_t*)p, x);+#elif defined(FUTHARK_OPENCL_DEVICE_ATOMICS)+  return atomic_fetch_or_explicit((volatile atomic_int*)p, x,+                                 memory_order_acq_rel,+                                 memory_scope_device); #else   return atomic_or(p, x); #endif@@ -219,6 +256,10 @@ SCALAR_FUN_ATTR int32_t atomic_xor_i32_global(volatile __global int32_t *p, int32_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicXor((int32_t*)p, x);+#elif defined(FUTHARK_OPENCL_DEVICE_ATOMICS)+  return atomic_fetch_xor_explicit((volatile atomic_int*)p, x,+                                  memory_order_acq_rel,+                                  memory_scope_device); #else   return atomic_xor(p, x); #endif
rts/c/atomics64.h view
@@ -33,6 +33,10 @@ SCALAR_FUN_ATTR int64_t atomic_xchg_i64_global(volatile __global int64_t *p, int64_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicExch((unsigned long long*)p, x);+#elif defined(FUTHARK_OPENCL_DEVICE_ATOMICS)+  return atomic_exchange_explicit((volatile atomic_long*)p, x,+                                  memory_order_acq_rel,+                                  memory_scope_device); #else   return atom_xchg(p, x); #endif@@ -50,6 +54,11 @@                                                          int64_t cmp, int64_t val) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicCAS((unsigned long long*)p, cmp, val);+#elif defined(FUTHARK_OPENCL_DEVICE_ATOMICS)+  atomic_compare_exchange_strong_explicit(+    (volatile atomic_long*)p, &cmp, val,+    memory_order_acq_rel, memory_order_acquire, memory_scope_device);+  return cmp; #else   return atom_cmpxchg(p, cmp, val); #endif@@ -67,6 +76,10 @@ SCALAR_FUN_ATTR int64_t atomic_add_i64_global(volatile __global int64_t *p, int64_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicAdd((unsigned long long*)p, x);+#elif defined(FUTHARK_OPENCL_DEVICE_ATOMICS)+  return atomic_fetch_add_explicit((volatile atomic_long*)p, x,+                                   memory_order_acq_rel,+                                   memory_scope_device); #else   return atom_add(p, x); #endif@@ -144,6 +157,10 @@     old = atomic_cmpxchg_i64_global((volatile __global int64_t*)p, assumed, old);   } while (assumed != old);   return old;+#elif defined(FUTHARK_OPENCL_DEVICE_ATOMICS)+  return atomic_fetch_max_explicit((volatile atomic_long*)p, x,+                                   memory_order_acq_rel,+                                   memory_scope_device); #else   return atom_max(p, x); #endif@@ -178,6 +195,10 @@     old = atomic_cmpxchg_i64_global((volatile __global int64_t*)p, assumed, old);   } while (assumed != old);   return old;+#elif defined(FUTHARK_OPENCL_DEVICE_ATOMICS)+  return atomic_fetch_min_explicit((volatile atomic_long*)p, x,+                                   memory_order_acq_rel,+                                   memory_scope_device); #else   return atom_min(p, x); #endif@@ -203,6 +224,10 @@ SCALAR_FUN_ATTR uint64_t atomic_umax_i64_global(volatile __global uint64_t *p, uint64_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicMax((unsigned long long*)p, x);+#elif defined(FUTHARK_OPENCL_DEVICE_ATOMICS)+  return atomic_fetch_max_explicit((volatile atomic_ulong*)p, x,+                                   memory_order_acq_rel,+                                   memory_scope_device); #else   return atom_max(p, x); #endif@@ -219,6 +244,10 @@ SCALAR_FUN_ATTR uint64_t atomic_umin_i64_global(volatile __global uint64_t *p, uint64_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicMin((unsigned long long*)p, x);+#elif defined(FUTHARK_OPENCL_DEVICE_ATOMICS)+  return atomic_fetch_min_explicit((volatile atomic_ulong*)p, x,+                                   memory_order_acq_rel,+                                   memory_scope_device); #else   return atom_min(p, x); #endif@@ -235,6 +264,10 @@ SCALAR_FUN_ATTR int64_t atomic_and_i64_global(volatile __global int64_t *p, int64_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicAnd((unsigned long long*)p, x);+#elif defined(FUTHARK_OPENCL_DEVICE_ATOMICS)+  return atomic_fetch_and_explicit((volatile atomic_long*)p, x,+                                   memory_order_acq_rel,+                                   memory_scope_device); #else   return atom_and(p, x); #endif@@ -251,6 +284,10 @@ SCALAR_FUN_ATTR int64_t atomic_or_i64_global(volatile __global int64_t *p, int64_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicOr((unsigned long long*)p, x);+#elif defined(FUTHARK_OPENCL_DEVICE_ATOMICS)+  return atomic_fetch_or_explicit((volatile atomic_long*)p, x,+                                  memory_order_acq_rel,+                                  memory_scope_device); #else   return atom_or(p, x); #endif@@ -267,6 +304,10 @@ SCALAR_FUN_ATTR int64_t atomic_xor_i64_global(volatile __global int64_t *p, int64_t x) { #if defined(FUTHARK_CUDA) || defined(FUTHARK_HIP)   return atomicXor((unsigned long long*)p, x);+#elif defined(FUTHARK_OPENCL_DEVICE_ATOMICS)+  return atomic_fetch_xor_explicit((volatile atomic_long*)p, x,+                                   memory_order_acq_rel,+                                   memory_scope_device); #else   return atom_xor(p, x); #endif
rts/c/backends/opencl.h view
@@ -508,6 +508,7 @@                              const char *extra_build_opts[],                              struct opencl_device_option device_option) {   int compile_opts_size = 1024;+  bool cl_std_was_set = false;    for (int i = 0; i < NUM_TUNING_PARAMS; i++) {     compile_opts_size += 2*(strlen(ctx->cfg->tuning_params[i].name) + 20);@@ -519,6 +520,7 @@    for (int i = 0; extra_build_opts[i] != NULL; i++) {     compile_opts_size += strlen(extra_build_opts[i] + 1);+    cl_std_was_set |= strncmp(extra_build_opts[i], "-cl-std=", 8) == 0;   }    for (int i = 0; i < num_macros; i++) {@@ -530,6 +532,14 @@   int w = snprintf(compile_opts, compile_opts_size,                    "-DLOCKSTEP_WIDTH=%d ",                    (int)ctx->lockstep_width);++  bool is_rusticl_asahi =+    strcmp(device_option.platform_name, "rusticl") == 0 &&+    strncmp(device_option.device_name, "Apple M", 7) == 0;+  if (is_rusticl_asahi && !cl_std_was_set) {+    w += snprintf(compile_opts+w, compile_opts_size-w,+                  "-cl-std=CL2.0 ");+  }    w += snprintf(compile_opts+w, compile_opts_size-w,                 "-D%s=%d ",
rts/c/scheduler.h view
@@ -101,6 +101,10 @@ #include <sys/sysinfo.h> #include <sys/resource.h> #include <signal.h>+#elif defined(__OpenBSD__)+#include <sys/resource.h>+#include <sys/sysctl.h>+#include <signal.h> #elif defined(__EMSCRIPTEN__) #include <emscripten/threading.h> #include <sys/sysinfo.h>@@ -157,6 +161,15 @@   return ncores; #elif defined(__linux__)   return get_nprocs();+#elif defined(__OpenBSD__)+  int mib[2], ncores;+  size_t len;++  mib[0] = CTL_HW;+  mib[1] = HW_NCPUONLINE;+  len = sizeof(ncores);+  CHECK_ERRNO(sysctl(mib, 2, &ncores, &len, NULL, 0), "sysctl");+  return ncores; #elif __EMSCRIPTEN__   return emscripten_num_logical_cores(); #else
rts/opencl/prelude.cl view
@@ -29,9 +29,21 @@ #pragma OPENCL EXTENSION cl_khr_int64_base_atomics : enable #pragma OPENCL EXTENSION cl_khr_int64_extended_atomics : enable +#if (__OPENCL_C_VERSION__ >= 200 && __OPENCL_C_VERSION__ < 300) || \+  (defined(__opencl_c_atomic_order_acq_rel) && \+   defined(__opencl_c_atomic_scope_device))+#define FUTHARK_OPENCL_DEVICE_ATOMICS+#endif+ // NVIDIAs OpenCL does not create device-wide memory fences (see #734), so we-// use inline assembly if we detect we are on an NVIDIA GPU.-#ifdef cl_nv_pragma_unroll+// use inline assembly if we detect we are on an NVIDIA GPU.  OpenCL 2.0+// provides a portable device-wide fence.+#ifdef FUTHARK_OPENCL_DEVICE_ATOMICS+static inline void mem_fence_global() {+  atomic_work_item_fence(CLK_LOCAL_MEM_FENCE | CLK_GLOBAL_MEM_FENCE,+                         memory_order_acq_rel, memory_scope_device);+}+#elif defined(cl_nv_pragma_unroll) static inline void mem_fence_global() {   asm("membar.gl;"); }
rts/python/opencl.py view
@@ -288,6 +288,15 @@     # compiler should provide us with the variables to which     # parameters are mapped.     if len(program_src) >= 0:+        is_rusticl_asahi = (+            self.platform.name == "rusticl"+            and self.device.name.startswith("Apple M")+        )+        if is_rusticl_asahi and not any(+            opt.startswith("-cl-std=") for opt in build_options+        ):+            build_options += ["-cl-std=CL2.0"]+         build_options += ["-DLOCKSTEP_WIDTH={}".format(lockstep_width)]          build_options += [
src-testing/Futhark/IR/Syntax/CoreTests.hs view
@@ -4,7 +4,6 @@  import Control.Applicative import Data.Loc (Loc (..), Pos (..))-import Futhark.IR.Pretty (prettyString) import Futhark.IR.Syntax.Core import Language.Futhark.CoreTests () import Language.Futhark.PrimitiveTests ()@@ -34,39 +33,6 @@     where       intconst = Constant . IntValue . Int32Value -subShapeTests :: [TestTree]-subShapeTests =-  [ shape [free 1, free 2] `isSubShapeOf` shape [free 1, free 2],-    shape [free 1, free 3] `isNotSubShapeOf` shape [free 1, free 2],-    shape [free 1] `isNotSubShapeOf` shape [free 1, free 2],-    shape [free 1, free 2] `isSubShapeOf` shape [free 1, Ext 3],-    shape [Ext 1, Ext 2] `isNotSubShapeOf` shape [Ext 1, Ext 1],-    shape [Ext 1, Ext 1] `isSubShapeOf` shape [Ext 1, Ext 2]-  ]-  where-    shape :: [ExtSize] -> ExtShape-    shape = Shape--    free :: Int -> ExtSize-    free = Free . Constant . IntValue . Int32Value . fromIntegral--    isSubShapeOf shape1 shape2 =-      subShapeTest shape1 shape2 True-    isNotSubShapeOf shape1 shape2 =-      subShapeTest shape1 shape2 False--    subShapeTest :: ExtShape -> ExtShape -> Bool -> TestTree-    subShapeTest shape1 shape2 expected =-      testCase-        ( "subshapeOf "-            ++ prettyString shape1-            ++ " "-            ++ prettyString shape2-            ++ " == "-            ++ show expected-        )-        $ shape1 `subShapeOf` shape2 @?= expected- provenanceTests :: [TestTree] provenanceTests =   [ testGroup@@ -95,6 +61,5 @@ tests =   testGroup     "Internal CoreTests"-    [ testGroup "subShape" subShapeTests,-      testGroup "Provenance" provenanceTests+    [ testGroup "Provenance" provenanceTests     ]
+ src-testing/Generated/AllFutBenchmarks.hs view
@@ -0,0 +1,24 @@+module Generated.AllFutBenchmarks+  ( allFutBenchmarkCases,+    BenchmarkCaseData,+  )+where++import Generated.AllFutBenchmarks.Accelerate.Nbody.Nbodybh qualified as AllFutBenchmarksAccelerateNbodyNbodybh+import Language.Futhark.TypeChecker.Constraints (CtTy, TyParams, TyVars)++type BenchmarkCaseData = ([CtTy ()], TyParams, TyVars ())++allFutBenchmarkCases :: [(String, BenchmarkCaseData)]+allFutBenchmarkCases =+  [ ("accelerate/nbody/nbody-bh.fut (Block 1/10) (Cons: 121)", head AllFutBenchmarksAccelerateNbodyNbodybh.benchmarkDataList),+    ("accelerate/nbody/nbody-bh.fut (Block 2/10) (Cons: 146)", AllFutBenchmarksAccelerateNbodyNbodybh.benchmarkDataList !! 1),+    ("accelerate/nbody/nbody-bh.fut (Block 3/10) (Cons: 133)", AllFutBenchmarksAccelerateNbodyNbodybh.benchmarkDataList !! 2),+    ("accelerate/nbody/nbody-bh.fut (Block 4/10) (Cons: 45)", AllFutBenchmarksAccelerateNbodyNbodybh.benchmarkDataList !! 3),+    ("accelerate/nbody/nbody-bh.fut (Block 5/10) (Cons: 210)", AllFutBenchmarksAccelerateNbodyNbodybh.benchmarkDataList !! 4),+    ("accelerate/nbody/nbody-bh.fut (Block 6/10) (Cons: 401)", AllFutBenchmarksAccelerateNbodyNbodybh.benchmarkDataList !! 5),+    ("accelerate/nbody/nbody-bh.fut (Block 7/10) (Cons: 39)", AllFutBenchmarksAccelerateNbodyNbodybh.benchmarkDataList !! 6),+    ("accelerate/nbody/nbody-bh.fut (Block 8/10) (Cons: 173)", AllFutBenchmarksAccelerateNbodyNbodybh.benchmarkDataList !! 7),+    ("accelerate/nbody/nbody-bh.fut (Block 9/10) (Cons: 164)", AllFutBenchmarksAccelerateNbodyNbodybh.benchmarkDataList !! 8),+    ("accelerate/nbody/nbody-bh.fut (Block 10/10) (Cons: 40)", AllFutBenchmarksAccelerateNbodyNbodybh.benchmarkDataList !! 9)+  ]
+ src-testing/Generated/AllFutBenchmarks/Accelerate/Nbody/Nbodybh.hs view
@@ -0,0 +1,1534 @@+module Generated.AllFutBenchmarks.Accelerate.Nbody.Nbodybh (benchmarkDataList) where++import Data.Map qualified as M+import Futhark.Util.Loc (Loc (NoLoc))+import Language.Futhark.Syntax+import Language.Futhark.SyntaxTests ()+import Language.Futhark.TypeChecker.Constraints+  ( CtTy (..),+    Reason (..),+    TyParams,+    TyVarInfo (..),+    TyVars,+  )++(~) :: TypeBase () NoUniqueness -> TypeBase () NoUniqueness -> CtTy ()+t1 ~ t2 = CtEq (Reason mempty) t1 t2++type BenchmarkCaseData = ([CtTy ()], TyParams, TyVars ())++benchmarkDataList :: [BenchmarkCaseData]+benchmarkDataList =+  [ ( [ "t_8322_8328_8326" ~ "[]t_8322_8328_8326_8322_8329_8327",+        "t_8322_8328_8328" ~ "[]t_8322_8328_8328_8322_8329_8328",+        "t_8321_8323_8326" ~ "[]t_8321_8323_8326_8322_8329_8329",+        "t_8327_8327" ~ "[]t_8327_8327_8323_8320_8320",+        "t_8321_8322_8321" ~ "[]t_8321_8322_8321_8323_8320_8321",+        "b_8327_8329" ~ "[]b_8327_8329_8323_8320_8322",+        "a_8327_8328" ~ "[]a_8327_8328_8323_8320_8323",+        "b_8326_8328" ~ "[]b_8326_8328_8323_8320_8324",+        "a_8326_8327" ~ "[]a_8326_8327_8323_8320_8325",+        "i32" ~ "t_8323",+        "num_8324" ~ "t_8323",+        "t_8323" ~ "i32",+        "t_8320" ~ "t_1",+        "i32" ~ "t_8322",+        "num_8321_8323" ~ "t_8322",+        "i32" ~ "i32",+        "t_8320" ~ "t_1",+        "t_8322" ~ "t_8321",+        "i32" ~ "t_8321",+        "t_8322_8326" ~ "arg_8323_8325 -> res_8323_8326",+        "t_8322_8327" ~ "arg_8323_8325",+        "res_8323_8326" ~ "arg_8323_8329 -> res_8324_8320",+        "t_8323_8321" ~ "arg_8323_8329",+        "t_8322_8326" ~ "arg_8324_8323 -> res_8324_8324",+        "t_8322_8328" ~ "arg_8324_8323",+        "res_8324_8324" ~ "arg_8324_8327 -> res_8324_8328",+        "t_8323_8322" ~ "arg_8324_8327",+        "t_8322_8326" ~ "arg_8325_8321 -> res_8325_8322",+        "t_8322_8329" ~ "arg_8325_8321",+        "res_8325_8322" ~ "arg_8325_8325 -> res_8325_8326",+        "t_8323_8323" ~ "arg_8325_8325",+        "t_8322_8326" ~ "arg_8325_8329 -> res_8326_8320",+        "t_8323_8320" ~ "arg_8325_8329",+        "res_8326_8320" ~ "arg_8326_8323 -> res_8326_8324",+        "t_8323_8324" ~ "arg_8326_8323",+        "{x: t_8320} -> t_8321" ~ "a_8326_8329 -> x_8327_8320",+        "[]t_1" ~ "a_8326_8327",+        "{as: []a_8326_8329} -> *[]x_8327_8320" ~ "a_8326_8327 -> b_8326_8328",+        "t_8327_8327" ~ "b_8326_8328",+        "t_8328_8322" ~ "t_8328_8323",+        "num_8328_8324" ~ "t_8328_8323",+        "bool" ~ "bool",+        "t_8328_8322" ~ "t_8329_8321",+        "num_8329_8322" ~ "t_8329_8321",+        "bool" ~ "bool",+        "t_8328_8322" ~ "t_8329_8329",+        "num_8321_8320_8320" ~ "t_8329_8329",+        "bool" ~ "bool",+        "t_8328_8322" ~ "t_8321_8320_8327",+        "num_8321_8320_8328" ~ "t_8321_8320_8327",+        "bool" ~ "bool",+        "{x: t_8328_8322} -> (i64, i64, i64, i64)" ~ "a_8328_8320 -> x_8328_8321",+        "t_8327_8327" ~ "a_8327_8328",+        "{as: []a_8328_8320} -> *[]x_8328_8321" ~ "a_8327_8328 -> b_8327_8329",+        "t_8321_8322_8321" ~ "b_8327_8329",+        "t_8321_8322_8323 -> t_8321_8322_8323 -> t_8321_8322_8323" ~ "t_8322_8326",+        "(t_8322_8327, t_8322_8328, t_8322_8329, t_8323_8320) -> (t_8323_8321, t_8323_8322, t_8323_8323, t_8323_8324) -> (res_8324_8320, res_8324_8328, res_8325_8326, res_8326_8324)" ~ "a_8321_8322_8322 -> a_8321_8322_8322 -> a_8321_8322_8322",+        "(num_8321_8322_8326, num_8321_8322_8327, num_8321_8322_8328, num_8321_8322_8329)" ~ "a_8321_8322_8322",+        "t_8321_8322_8321" ~ "[]a_8321_8322_8322",+        "t_8321_8323_8326" ~ "[]a_8321_8322_8322",+        "t_8321_8323_8326" ~ "[]t_8321_8323_8327",+        "(t_8321_8324_8320, t_8321_8324_8321, t_8321_8324_8322, t_8321_8324_8323)" ~ "t_8321_8323_8327",+        "t_8321_8325_8327" ~ "num_8321_8325_8326",+        "t_8321_8324_8324" ~ "t_8321_8325_8329",+        "num_8321_8326_8320" ~ "t_8321_8325_8329",+        "bool" ~ "bool",+        "t_8321_8324_8325" ~ "t_8321_8325_8328",+        "i64" ~ "t_8321_8325_8328",+        "num_8321_8325_8326" ~ "t_8321_8325_8325",+        "t_8321_8325_8328" ~ "t_8321_8325_8325",+        "t_8321_8324_8324" ~ "t_8321_8327_8326",+        "num_8321_8327_8327" ~ "t_8321_8327_8326",+        "bool" ~ "bool",+        "t_8321_8324_8320" ~ "t_8321_8327_8325",+        "i64" ~ "t_8321_8327_8325",+        "t_8321_8325_8325" ~ "t_8321_8325_8324",+        "t_8321_8327_8325" ~ "t_8321_8325_8324",+        "t_8321_8324_8324" ~ "t_8321_8329_8323",+        "num_8321_8329_8324" ~ "t_8321_8329_8323",+        "bool" ~ "bool",+        "t_8321_8324_8326" ~ "t_8321_8329_8322",+        "i64" ~ "t_8321_8329_8322",+        "t_8321_8325_8324" ~ "t_8321_8325_8323",+        "t_8321_8329_8322" ~ "t_8321_8325_8323",+        "t_8321_8324_8324" ~ "t_8322_8321_8320",+        "num_8322_8321_8321" ~ "t_8322_8321_8320",+        "bool" ~ "bool",+        "t_8321_8324_8321" ~ "t_8322_8320_8329",+        "i64" ~ "t_8322_8320_8329",+        "t_8321_8325_8323" ~ "t_8321_8325_8322",+        "t_8322_8320_8329" ~ "t_8321_8325_8322",+        "t_8321_8324_8324" ~ "t_8322_8322_8327",+        "num_8322_8322_8328" ~ "t_8322_8322_8327",+        "bool" ~ "bool",+        "t_8321_8324_8327" ~ "t_8322_8322_8326",+        "i64" ~ "t_8322_8322_8326",+        "t_8321_8325_8322" ~ "t_8321_8325_8321",+        "t_8322_8322_8326" ~ "t_8321_8325_8321",+        "t_8321_8324_8324" ~ "t_8322_8324_8324",+        "num_8322_8324_8325" ~ "t_8322_8324_8324",+        "bool" ~ "bool",+        "t_8321_8324_8322" ~ "t_8322_8324_8323",+        "i64" ~ "t_8322_8324_8323",+        "t_8321_8325_8321" ~ "t_8321_8325_8320",+        "t_8322_8324_8323" ~ "t_8321_8325_8320",+        "t_8321_8324_8324" ~ "t_8322_8326_8321",+        "num_8322_8326_8322" ~ "t_8322_8326_8321",+        "bool" ~ "bool",+        "t_8321_8324_8328" ~ "t_8322_8326_8320",+        "i64" ~ "t_8322_8326_8320",+        "t_8321_8325_8320" ~ "t_8321_8324_8329",+        "t_8322_8326_8320" ~ "t_8321_8324_8329",+        "{bin: t_8321_8324_8324} -> (t_8321_8324_8325, t_8321_8324_8326, t_8321_8324_8327, t_8321_8324_8328) -> t_8321_8324_8329" ~ "a_8322_8327_8327 -> b_8322_8327_8328 -> x_8322_8327_8329",+        "t_8327_8327" ~ "[]a_8322_8327_8327",+        "t_8321_8323_8326" ~ "[]b_8322_8327_8328",+        "t_8322_8328_8326" ~ "[]x_8322_8327_8329",+        "[]t_1" ~ "t_8322_8328_8328",+        "t_8322_8328_8328" ~ "[]t_8322_8328_8327",+        "t_8322_8328_8326" ~ "[]i64",+        "[]t_1" ~ "[]t_8322_8328_8327",+        "[]t_1" ~ "[]t_8322_8328_8327"+      ],+      M.fromList [("t_1", (0, Unlifted, NoLoc))],+      M.fromList [("t_8320", (5, TyVarFree NoLoc Lifted)), ("t_8321", (6, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322", (6, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323", (6, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8324", (6, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8323", (6, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8326", (5, TyVarFree NoLoc Lifted)), ("t_8322_8327", (6, TyVarFree NoLoc Lifted)), ("t_8322_8328", (6, TyVarFree NoLoc Lifted)), ("t_8322_8329", (6, TyVarFree NoLoc Lifted)), ("t_8323_8320", (6, TyVarFree NoLoc Lifted)), ("t_8323_8321", (7, TyVarFree NoLoc Lifted)), ("t_8323_8322", (7, TyVarFree NoLoc Lifted)), ("t_8323_8323", (7, TyVarFree NoLoc Lifted)), ("t_8323_8324", (7, TyVarFree NoLoc Lifted)), ("arg_8323_8325", (8, TyVarFree NoLoc Lifted)), ("res_8323_8326", (8, TyVarFree NoLoc Lifted)), ("arg_8323_8329", (8, TyVarFree NoLoc Lifted)), ("res_8324_8320", (8, TyVarFree NoLoc Lifted)), ("arg_8324_8323", (8, TyVarFree NoLoc Lifted)), ("res_8324_8324", (8, TyVarFree NoLoc Lifted)), ("arg_8324_8327", (8, TyVarFree NoLoc Lifted)), ("res_8324_8328", (8, TyVarFree NoLoc Lifted)), ("arg_8325_8321", (8, TyVarFree NoLoc Lifted)), ("res_8325_8322", (8, TyVarFree NoLoc Lifted)), ("arg_8325_8325", (8, TyVarFree NoLoc Lifted)), ("res_8325_8326", (8, TyVarFree NoLoc Lifted)), ("arg_8325_8329", (8, TyVarFree NoLoc Lifted)), ("res_8326_8320", (8, TyVarFree NoLoc Lifted)), ("arg_8326_8323", (8, TyVarFree NoLoc Lifted)), ("res_8326_8324", (8, TyVarFree NoLoc Lifted)), ("a_8326_8327", (4, TyVarFree NoLoc Lifted)), ("b_8326_8328", (4, TyVarFree NoLoc Lifted)), ("a_8326_8329", (4, TyVarFree NoLoc Unlifted)), ("x_8327_8320", (4, TyVarFree NoLoc Unlifted)), ("t_8327_8327", (5, TyVarFree NoLoc Lifted)), ("a_8327_8328", (6, TyVarFree NoLoc Lifted)), ("b_8327_8329", (6, TyVarFree NoLoc Lifted)), ("a_8328_8320", (6, TyVarFree NoLoc Unlifted)), ("x_8328_8321", (6, TyVarFree NoLoc Unlifted)), ("t_8328_8322", (7, TyVarFree NoLoc Lifted)), ("t_8328_8323", (8, TyVarFree NoLoc Unlifted)), ("num_8328_8324", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8329_8321", (8, TyVarFree NoLoc Unlifted)), ("num_8329_8322", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8329_8329", (8, TyVarFree NoLoc Unlifted)), ("num_8321_8320_8320", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8320_8327", (8, TyVarFree NoLoc Unlifted)), ("num_8321_8320_8328", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8322_8321", (7, TyVarFree NoLoc Lifted)), ("a_8321_8322_8322", (8, TyVarFree NoLoc Unlifted)), ("t_8321_8322_8323", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8322_8326", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8322_8327", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8322_8328", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8322_8329", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8323_8326", (9, TyVarFree NoLoc Lifted)), ("t_8321_8323_8327", (10, TyVarFree NoLoc Unlifted)), ("t_8321_8324_8320", (11, TyVarFree NoLoc Lifted)), ("t_8321_8324_8321", (11, TyVarFree NoLoc Lifted)), ("t_8321_8324_8322", (11, TyVarFree NoLoc Lifted)), ("t_8321_8324_8323", (11, TyVarFree NoLoc Lifted)), ("t_8321_8324_8324", (13, TyVarFree NoLoc Lifted)), ("t_8321_8324_8325", (14, TyVarFree NoLoc Lifted)), ("t_8321_8324_8326", (14, TyVarFree NoLoc Lifted)), ("t_8321_8324_8327", (14, TyVarFree NoLoc Lifted)), ("t_8321_8324_8328", (14, TyVarFree NoLoc Lifted)), ("t_8321_8324_8329", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8325_8320", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8325_8321", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8325_8322", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8325_8323", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8325_8324", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8325_8325", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8325_8326", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8325_8327", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8325_8328", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8325_8329", (15, TyVarFree NoLoc Unlifted)), ("num_8321_8326_8320", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8327_8325", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8327_8326", (15, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8327_8327", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8329_8322", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8329_8323", (15, TyVarFree NoLoc Unlifted)), ("num_8321_8329_8324", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8320_8329", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8321_8320", (15, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8321_8321", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8322_8326", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8322_8327", (15, TyVarFree NoLoc Unlifted)), ("num_8322_8322_8328", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8324_8323", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8324_8324", (15, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8324_8325", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8326_8320", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8326_8321", (15, TyVarFree NoLoc Unlifted)), ("num_8322_8326_8322", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("a_8322_8327_8327", (12, TyVarFree NoLoc Unlifted)), ("b_8322_8327_8328", (12, TyVarFree NoLoc Unlifted)), ("x_8322_8327_8329", (12, TyVarFree NoLoc Unlifted)), ("t_8322_8328_8326", (13, TyVarFree NoLoc Lifted)), ("t_8322_8328_8327", (14, TyVarFree NoLoc Unlifted)), ("t_8322_8328_8328", (14, TyVarFree NoLoc Unlifted)), ("t_8322_8328_8326_8322_8329_8327", (13, TyVarFree NoLoc Lifted)), ("t_8322_8328_8328_8322_8329_8328", (14, TyVarFree NoLoc Unlifted)), ("t_8321_8323_8326_8322_8329_8329", (9, TyVarFree NoLoc Lifted)), ("t_8327_8327_8323_8320_8320", (5, TyVarFree NoLoc Lifted)), ("t_8321_8322_8321_8323_8320_8321", (7, TyVarFree NoLoc Lifted)), ("b_8327_8329_8323_8320_8322", (6, TyVarFree NoLoc Lifted)), ("a_8327_8328_8323_8320_8323", (6, TyVarFree NoLoc Lifted)), ("b_8326_8328_8323_8320_8324", (4, TyVarFree NoLoc Lifted)), ("a_8326_8327_8323_8320_8325", (4, TyVarFree NoLoc Lifted))]+    ),+    ( [ "t_8323_8320_8321" ~ "[]t_8323_8320_8321_8323_8323_8328",+        "t_8323_8320_8323" ~ "[]t_8323_8320_8323_8323_8323_8329",+        "t_8321_8323_8328" ~ "[]t_8321_8323_8328_8323_8324_8320",+        "t_8327_8329" ~ "[]t_8327_8329_8323_8324_8321",+        "t_8321_8322_8323" ~ "[]t_8321_8322_8323_8323_8324_8322",+        "b_8328_8321" ~ "[]b_8328_8321_8323_8324_8323",+        "a_8328_8320" ~ "[]a_8328_8320_8323_8324_8324",+        "b_8327_8320" ~ "[]b_8327_8320_8323_8324_8325",+        "a_8326_8329" ~ "[]a_8326_8329_8323_8324_8326",+        "i32" ~ "t_8323",+        "num_8324" ~ "t_8323",+        "t_8323" ~ "i32",+        "t_8320" ~ "t_1",+        "i32" ~ "t_8322",+        "num_8321_8323" ~ "t_8322",+        "i32" ~ "i32",+        "t_8320" ~ "t_1",+        "t_8322" ~ "t_8321",+        "i32" ~ "t_8321",+        "t_8321" ~ "i32",+        "t_8322_8328" ~ "arg_8323_8327 -> res_8323_8328",+        "t_8322_8329" ~ "arg_8323_8327",+        "res_8323_8328" ~ "arg_8324_8321 -> res_8324_8322",+        "t_8323_8323" ~ "arg_8324_8321",+        "t_8322_8328" ~ "arg_8324_8325 -> res_8324_8326",+        "t_8323_8320" ~ "arg_8324_8325",+        "res_8324_8326" ~ "arg_8324_8329 -> res_8325_8320",+        "t_8323_8324" ~ "arg_8324_8329",+        "t_8322_8328" ~ "arg_8325_8323 -> res_8325_8324",+        "t_8323_8321" ~ "arg_8325_8323",+        "res_8325_8324" ~ "arg_8325_8327 -> res_8325_8328",+        "t_8323_8325" ~ "arg_8325_8327",+        "t_8322_8328" ~ "arg_8326_8321 -> res_8326_8322",+        "t_8323_8322" ~ "arg_8326_8321",+        "res_8326_8322" ~ "arg_8326_8325 -> res_8326_8326",+        "t_8323_8326" ~ "arg_8326_8325",+        "{x: t_8320} -> i16" ~ "a_8327_8321 -> x_8327_8322",+        "[]t_1" ~ "a_8326_8329",+        "{as: []a_8327_8321} -> *[]x_8327_8322" ~ "a_8326_8329 -> b_8327_8320",+        "t_8327_8329" ~ "b_8327_8320",+        "t_8328_8324" ~ "t_8328_8325",+        "num_8328_8326" ~ "t_8328_8325",+        "bool" ~ "bool",+        "t_8328_8324" ~ "t_8329_8323",+        "num_8329_8324" ~ "t_8329_8323",+        "bool" ~ "bool",+        "t_8328_8324" ~ "t_8321_8320_8321",+        "num_8321_8320_8322" ~ "t_8321_8320_8321",+        "bool" ~ "bool",+        "t_8328_8324" ~ "t_8321_8320_8329",+        "num_8321_8321_8320" ~ "t_8321_8320_8329",+        "bool" ~ "bool",+        "{x: t_8328_8324} -> (i16, i16, i16, i16)" ~ "a_8328_8322 -> x_8328_8323",+        "t_8327_8329" ~ "a_8328_8320",+        "{as: []a_8328_8322} -> *[]x_8328_8323" ~ "a_8328_8320 -> b_8328_8321",+        "t_8321_8322_8323" ~ "b_8328_8321",+        "t_8321_8322_8325 -> t_8321_8322_8325 -> t_8321_8322_8325" ~ "t_8322_8328",+        "(t_8322_8329, t_8323_8320, t_8323_8321, t_8323_8322) -> (t_8323_8323, t_8323_8324, t_8323_8325, t_8323_8326) -> (res_8324_8322, res_8325_8320, res_8325_8328, res_8326_8326)" ~ "a_8321_8322_8324 -> a_8321_8322_8324 -> a_8321_8322_8324",+        "(num_8321_8322_8328, num_8321_8322_8329, num_8321_8323_8320, num_8321_8323_8321)" ~ "a_8321_8322_8324",+        "t_8321_8322_8323" ~ "[]a_8321_8322_8324",+        "t_8321_8323_8328" ~ "[]a_8321_8322_8324",+        "i64" ~ "t_8321_8323_8329",+        "num_8321_8324_8320" ~ "t_8321_8323_8329",+        "t_8321_8323_8328" ~ "[]t_8321_8324_8329",+        "bool" ~ "bool",+        "(num_8321_8324_8325, num_8321_8324_8326, num_8321_8324_8327, num_8321_8324_8328)" ~ "if_t_8321_8325_8322",+        "t_8321_8324_8329" ~ "if_t_8321_8325_8322",+        "(t_8321_8325_8323, t_8321_8325_8324, t_8321_8325_8325, t_8321_8325_8326)" ~ "if_t_8321_8325_8322",+        "t_8321_8327_8320" ~ "num_8321_8326_8329",+        "t_8321_8325_8327" ~ "t_8321_8327_8322",+        "num_8321_8327_8323" ~ "t_8321_8327_8322",+        "bool" ~ "bool",+        "t_8321_8325_8328" ~ "t_8321_8327_8321",+        "i16" ~ "t_8321_8327_8321",+        "num_8321_8326_8329" ~ "t_8321_8326_8328",+        "t_8321_8327_8321" ~ "t_8321_8326_8328",+        "t_8321_8325_8327" ~ "t_8321_8328_8329",+        "num_8321_8329_8320" ~ "t_8321_8328_8329",+        "bool" ~ "bool",+        "t_8321_8325_8323" ~ "t_8321_8328_8328",+        "i16" ~ "t_8321_8328_8328",+        "t_8321_8326_8328" ~ "t_8321_8326_8327",+        "t_8321_8328_8328" ~ "t_8321_8326_8327",+        "t_8321_8325_8327" ~ "t_8322_8320_8326",+        "num_8322_8320_8327" ~ "t_8322_8320_8326",+        "bool" ~ "bool",+        "t_8321_8325_8329" ~ "t_8322_8320_8325",+        "i16" ~ "t_8322_8320_8325",+        "t_8321_8326_8327" ~ "t_8321_8326_8326",+        "t_8322_8320_8325" ~ "t_8321_8326_8326",+        "t_8321_8325_8327" ~ "t_8322_8322_8323",+        "num_8322_8322_8324" ~ "t_8322_8322_8323",+        "bool" ~ "bool",+        "t_8321_8325_8324" ~ "t_8322_8322_8322",+        "i16" ~ "t_8322_8322_8322",+        "t_8321_8326_8326" ~ "t_8321_8326_8325",+        "t_8322_8322_8322" ~ "t_8321_8326_8325",+        "t_8321_8325_8327" ~ "t_8322_8324_8320",+        "num_8322_8324_8321" ~ "t_8322_8324_8320",+        "bool" ~ "bool",+        "t_8321_8326_8320" ~ "t_8322_8323_8329",+        "i16" ~ "t_8322_8323_8329",+        "t_8321_8326_8325" ~ "t_8321_8326_8324",+        "t_8322_8323_8329" ~ "t_8321_8326_8324",+        "t_8321_8325_8327" ~ "t_8322_8325_8327",+        "num_8322_8325_8328" ~ "t_8322_8325_8327",+        "bool" ~ "bool",+        "t_8321_8325_8325" ~ "t_8322_8325_8326",+        "i16" ~ "t_8322_8325_8326",+        "t_8321_8326_8324" ~ "t_8321_8326_8323",+        "t_8322_8325_8326" ~ "t_8321_8326_8323",+        "t_8321_8325_8327" ~ "t_8322_8327_8324",+        "num_8322_8327_8325" ~ "t_8322_8327_8324",+        "bool" ~ "bool",+        "t_8321_8326_8321" ~ "t_8322_8327_8323",+        "i16" ~ "t_8322_8327_8323",+        "t_8321_8326_8323" ~ "t_8321_8326_8322",+        "t_8322_8327_8323" ~ "t_8321_8326_8322",+        "t_8321_8326_8322" ~ "i16",+        "{bin: t_8321_8325_8327} -> (t_8321_8325_8328, t_8321_8325_8329, t_8321_8326_8320, t_8321_8326_8321) -> i64" ~ "a_8322_8329_8322 -> b_8322_8329_8323 -> x_8322_8329_8324",+        "t_8327_8329" ~ "[]a_8322_8329_8322",+        "t_8321_8323_8328" ~ "[]b_8322_8329_8323",+        "t_8323_8320_8321" ~ "[]x_8322_8329_8324",+        "[]t_1" ~ "t_8323_8320_8323",+        "t_8323_8320_8323" ~ "[]t_8323_8320_8322",+        "t_8323_8320_8321" ~ "[]i64",+        "[]t_1" ~ "[]t_8323_8320_8322",+        "t_8321_8325_8323" ~ "et_8323_8321_8324",+        "t_8321_8325_8324" ~ "et_8323_8321_8324",+        "t_8321_8325_8325" ~ "et_8323_8321_8324",+        "t_8321_8325_8326" ~ "et_8323_8321_8324",+        "i16 -> i64" ~ "a_8323_8321_8322 -> x_8323_8321_8323",+        "[]et_8323_8321_8324" ~ "[]a_8323_8321_8322",+        "num_8323_8322_8320" ~ "et_8323_8321_8329",+        "t_8321_8325_8323" ~ "et_8323_8321_8329",+        "t_8321_8325_8323" ~ "t_8323_8322_8321",+        "t_8321_8325_8324" ~ "t_8323_8322_8321",+        "t_8323_8322_8321" ~ "et_8323_8321_8329",+        "t_8321_8325_8323" ~ "t_8323_8322_8327",+        "t_8321_8325_8324" ~ "t_8323_8322_8327",+        "t_8323_8322_8327" ~ "t_8323_8322_8326",+        "t_8321_8325_8325" ~ "t_8323_8322_8326",+        "t_8323_8322_8326" ~ "et_8323_8321_8329",+        "num_8323_8323_8326" ~ "i64",+        "num_8323_8323_8327" ~ "i64",+        "([]t_1, []i64, []i16)" ~ "([]t_8323_8320_8322, []x_8323_8321_8323, []et_8323_8321_8329)"+      ],+      M.fromList [("t_1", (0, Unlifted, NoLoc))],+      M.fromList [("t_8320", (5, TyVarFree NoLoc Lifted)), ("t_8321", (6, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322", (6, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323", (6, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8324", (6, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8323", (6, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8328", (5, TyVarFree NoLoc Lifted)), ("t_8322_8329", (6, TyVarFree NoLoc Lifted)), ("t_8323_8320", (6, TyVarFree NoLoc Lifted)), ("t_8323_8321", (6, TyVarFree NoLoc Lifted)), ("t_8323_8322", (6, TyVarFree NoLoc Lifted)), ("t_8323_8323", (7, TyVarFree NoLoc Lifted)), ("t_8323_8324", (7, TyVarFree NoLoc Lifted)), ("t_8323_8325", (7, TyVarFree NoLoc Lifted)), ("t_8323_8326", (7, TyVarFree NoLoc Lifted)), ("arg_8323_8327", (8, TyVarFree NoLoc Lifted)), ("res_8323_8328", (8, TyVarFree NoLoc Lifted)), ("arg_8324_8321", (8, TyVarFree NoLoc Lifted)), ("res_8324_8322", (8, TyVarFree NoLoc Lifted)), ("arg_8324_8325", (8, TyVarFree NoLoc Lifted)), ("res_8324_8326", (8, TyVarFree NoLoc Lifted)), ("arg_8324_8329", (8, TyVarFree NoLoc Lifted)), ("res_8325_8320", (8, TyVarFree NoLoc Lifted)), ("arg_8325_8323", (8, TyVarFree NoLoc Lifted)), ("res_8325_8324", (8, TyVarFree NoLoc Lifted)), ("arg_8325_8327", (8, TyVarFree NoLoc Lifted)), ("res_8325_8328", (8, TyVarFree NoLoc Lifted)), ("arg_8326_8321", (8, TyVarFree NoLoc Lifted)), ("res_8326_8322", (8, TyVarFree NoLoc Lifted)), ("arg_8326_8325", (8, TyVarFree NoLoc Lifted)), ("res_8326_8326", (8, TyVarFree NoLoc Lifted)), ("a_8326_8329", (4, TyVarFree NoLoc Lifted)), ("b_8327_8320", (4, TyVarFree NoLoc Lifted)), ("a_8327_8321", (4, TyVarFree NoLoc Unlifted)), ("x_8327_8322", (4, TyVarFree NoLoc Unlifted)), ("t_8327_8329", (5, TyVarFree NoLoc Lifted)), ("a_8328_8320", (6, TyVarFree NoLoc Lifted)), ("b_8328_8321", (6, TyVarFree NoLoc Lifted)), ("a_8328_8322", (6, TyVarFree NoLoc Unlifted)), ("x_8328_8323", (6, TyVarFree NoLoc Unlifted)), ("t_8328_8324", (7, TyVarFree NoLoc Lifted)), ("t_8328_8325", (8, TyVarFree NoLoc Unlifted)), ("num_8328_8326", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8329_8323", (8, TyVarFree NoLoc Unlifted)), ("num_8329_8324", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8320_8321", (8, TyVarFree NoLoc Unlifted)), ("num_8321_8320_8322", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8320_8329", (8, TyVarFree NoLoc Unlifted)), ("num_8321_8321_8320", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8322_8323", (7, TyVarFree NoLoc Lifted)), ("a_8321_8322_8324", (8, TyVarFree NoLoc Unlifted)), ("t_8321_8322_8325", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8322_8328", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8322_8329", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8323_8320", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8323_8321", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8323_8328", (9, TyVarFree NoLoc Lifted)), ("t_8321_8323_8329", (10, TyVarFree NoLoc Unlifted)), ("num_8321_8324_8320", (10, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8324_8325", (10, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8324_8326", (10, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8324_8327", (10, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8324_8328", (10, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8324_8329", (10, TyVarFree NoLoc Unlifted)), ("if_t_8321_8325_8322", (10, TyVarFree NoLoc SizeLifted)), ("t_8321_8325_8323", (11, TyVarFree NoLoc Lifted)), ("t_8321_8325_8324", (11, TyVarFree NoLoc Lifted)), ("t_8321_8325_8325", (11, TyVarFree NoLoc Lifted)), ("t_8321_8325_8326", (11, TyVarFree NoLoc Lifted)), ("t_8321_8325_8327", (13, TyVarFree NoLoc Lifted)), ("t_8321_8325_8328", (14, TyVarFree NoLoc Lifted)), ("t_8321_8325_8329", (14, TyVarFree NoLoc Lifted)), ("t_8321_8326_8320", (14, TyVarFree NoLoc Lifted)), ("t_8321_8326_8321", (14, TyVarFree NoLoc Lifted)), ("t_8321_8326_8322", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8326_8323", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8326_8324", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8326_8325", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8326_8326", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8326_8327", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8326_8328", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8326_8329", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8327_8320", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8327_8321", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8327_8322", (15, TyVarFree NoLoc Unlifted)), ("num_8321_8327_8323", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8328_8328", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8328_8329", (15, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8329_8320", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8320_8325", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8320_8326", (15, TyVarFree NoLoc Unlifted)), ("num_8322_8320_8327", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8322_8322", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8322_8323", (15, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8322_8324", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8323_8329", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8324_8320", (15, TyVarFree NoLoc Unlifted)), ("num_8322_8324_8321", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8325_8326", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8325_8327", (15, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8325_8328", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8327_8323", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8327_8324", (15, TyVarFree NoLoc Unlifted)), ("num_8322_8327_8325", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("a_8322_8329_8322", (12, TyVarFree NoLoc Unlifted)), ("b_8322_8329_8323", (12, TyVarFree NoLoc Unlifted)), ("x_8322_8329_8324", (12, TyVarFree NoLoc Unlifted)), ("t_8323_8320_8321", (13, TyVarFree NoLoc Lifted)), ("t_8323_8320_8322", (14, TyVarFree NoLoc Unlifted)), ("t_8323_8320_8323", (14, TyVarFree NoLoc Unlifted)), ("a_8323_8321_8322", (14, TyVarFree NoLoc Unlifted)), ("x_8323_8321_8323", (14, TyVarFree NoLoc Unlifted)), ("et_8323_8321_8324", (14, TyVarFree NoLoc Unlifted)), ("et_8323_8321_8329", (14, TyVarFree NoLoc Unlifted)), ("num_8323_8322_8320", (14, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8322_8321", (14, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8322_8326", (14, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8322_8327", (14, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8323_8326", (4, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8323_8327", (4, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8320_8321_8323_8323_8328", (13, TyVarFree NoLoc Lifted)), ("t_8323_8320_8323_8323_8323_8329", (14, TyVarFree NoLoc Unlifted)), ("t_8321_8323_8328_8323_8324_8320", (9, TyVarFree NoLoc Lifted)), ("t_8327_8329_8323_8324_8321", (5, TyVarFree NoLoc Lifted)), ("t_8321_8322_8323_8323_8324_8322", (7, TyVarFree NoLoc Lifted)), ("b_8328_8321_8323_8324_8323", (6, TyVarFree NoLoc Lifted)), ("a_8328_8320_8323_8324_8324", (6, TyVarFree NoLoc Lifted)), ("b_8327_8320_8323_8324_8325", (4, TyVarFree NoLoc Lifted)), ("a_8326_8329_8323_8324_8326", (4, TyVarFree NoLoc Lifted))]+    ),+    ( [ "t_8328_8325" ~ "[]t_8328_8325_8322_8321_8328",+        "t_8322_8321_8320" ~ "[]t_8322_8321_8320_8322_8321_8329",+        "index_8321_8327_8326" ~ "[]index_8321_8327_8326_8322_8322_8320",+        "t_8329_8322" ~ "[][]t_8329_8322_8322_8322_8321",+        "index_elem_8321_8327_8327" ~ "[]index_elem_8321_8327_8327_8322_8322_8322",+        "index_8321_8326_8329" ~ "[]index_8321_8326_8329_8322_8322_8323",+        "t_8321_8323_8324" ~ "[][]t_8321_8323_8324_8322_8322_8324",+        "index_elem_8321_8327_8320" ~ "[]index_elem_8321_8327_8320_8322_8322_8325",+        "b_8329_8324" ~ "[][]b_8329_8324_8322_8322_8326",+        "a_8329_8323" ~ "[][]a_8329_8323_8322_8322_8327",+        "b_8329_8326" ~ "[][]b_8329_8326_8322_8322_8328",+        "a_8329_8325" ~ "[]a_8329_8325_8322_8322_8329",+        "b_8329_8328" ~ "[]b_8329_8328_8322_8323_8320",+        "a_8329_8327" ~ "[]a_8329_8327_8322_8323_8321",+        "b_8321_8320_8320" ~ "[]b_8321_8320_8320_8322_8323_8322",+        "a_8329_8329" ~ "[][]a_8329_8329_8322_8323_8323",+        "b_8321_8320_8322" ~ "[][]b_8321_8320_8322_8322_8323_8324",+        "a_8321_8320_8321" ~ "[][]a_8321_8320_8321_8322_8323_8325",+        "t_8326_8321" ~ "[][]t_8326_8321_8322_8323_8326",+        "t_8328_8326" ~ "[]t_8328_8326_8322_8323_8327",+        "et_8328_8327" ~ "[]et_8328_8327_8322_8323_8328",+        "t_8325_8324" ~ "[][]t_8325_8324_8322_8323_8329",+        "t_8322_8324" ~ "[]t_8322_8324_8322_8324_8320",+        "t_8322_8322" ~ "[]t_8322_8322_8322_8324_8321",+        "t_8325_8322" ~ "[][]t_8325_8322_8322_8324_8322",+        "t_8325_8325" ~ "[]t_8325_8325_8322_8324_8323",+        "et_8325_8326" ~ "[]et_8325_8326_8322_8324_8324",+        "t_8325_8323" ~ "[][]t_8325_8323_8322_8324_8325",+        "t_8322_8323" ~ "[]t_8322_8323_8322_8324_8326",+        "a_8322_8325" ~ "[]a_8322_8325_8322_8324_8327",+        "b_8322_8328" ~ "[]b_8322_8328_8322_8324_8328",+        "a_8323_8322" ~ "[]a_8323_8322_8322_8324_8329",+        "a_8322_8327" ~ "[][]a_8322_8327_8322_8325_8320",+        "t_8321_8323" ~ "[]t_8321_8323_8322_8325_8321",+        "t_8321_8324" ~ "[]t_8321_8324_8322_8325_8322",+        "i64" ~ "t_8321",+        "i64" ~ "t_8321",+        "t_8321" ~ "t_8320",+        "i64" ~ "t_8320",+        "t_8320" ~ "i64",+        "[]t_1" ~ "[]t_8321_8320",+        "(t_8321_8323, t_8321_8324)" ~ "([]t_8321_8320, []t_8321_8320)",+        "i32 -> t_1 -> i32" ~ "i32 -> t_8321_8325 -> i32",+        "i32" ~ "i32",+        "t_8321_8324" ~ "[]t_8321_8325",+        "(t_8322_8322, t_8322_8323, t_8322_8324)" ~ "([]t_8321_8325, []i64, []i16)",+        "t_8321_8323" ~ "[]t_8322_8329",+        "i32 -> t_1 -> i32" ~ "i32 -> t_8323_8324 -> i32",+        "i32" ~ "i32",+        "{xs: []t_8323_8324} -> ([]t_8323_8324, []i64, []i16)" ~ "a_8323_8322 -> x_8323_8323",+        "[][]t_8322_8329" ~ "a_8322_8327",+        "{as: []a_8323_8322} -> *[]x_8323_8323" ~ "a_8322_8327 -> b_8322_8328",+        "b_8322_8328" ~ "a_8322_8325",+        "{xs: [](a_8324_8325, b_8324_8326, c_8324_8327)} -> ([]a_8324_8325, []b_8324_8326, []c_8324_8327)" ~ "a_8322_8325 -> b_8322_8326",+        "(t_8325_8322, t_8325_8323, t_8325_8324)" ~ "b_8322_8326",+        "t_8322_8323" ~ "et_8325_8326",+        "t_8325_8323" ~ "[]t_8325_8325",+        "[]et_8325_8326" ~ "[]t_8325_8325",+        "t_8326_8321" ~ "[]t_8325_8325",+        "i64" ~ "t_8326_8324",+        "i64" ~ "t_8326_8324",+        "t_8326_8324" ~ "t_8326_8323",+        "i64" ~ "t_8326_8323",+        "t_8325_8322" ~ "[][]t_8327_8324",+        "[]t_8327_8324" ~ "[]t_8327_8323",+        "t_8322_8322" ~ "[]t_8327_8323",+        "t_8326_8323" ~ "i64",+        "[]t_8327_8323" ~ "[]t_8326_8322",+        "t_8328_8325" ~ "[]t_8326_8322",+        "t_8322_8324" ~ "et_8328_8327",+        "t_8325_8324" ~ "[]t_8328_8326",+        "[]et_8328_8327" ~ "[]t_8328_8326",+        "t_8329_8322" ~ "[]t_8328_8326",+        "t_8326_8321" ~ "a_8321_8320_8321",+        "{a: [][]t_8321_8320_8323} -> [][]t_8321_8320_8323" ~ "a_8321_8320_8321 -> b_8321_8320_8322",+        "b_8321_8320_8322" ~ "a_8329_8329",+        "{xs: [][]t_8321_8320_8328} -> []t_8321_8320_8328" ~ "a_8329_8329 -> b_8321_8320_8320",+        "t_8321_8321_8324 -> t_8321_8321_8324 -> t_8321_8321_8324" ~ "update_elem_8321_8321_8323 -> update_elem_8321_8321_8323 -> update_elem_8321_8321_8323",+        "num_8321_8321_8325" ~ "update_elem_8321_8321_8323",+        "b_8321_8320_8320" ~ "a_8329_8327",+        "{xs: []update_elem_8321_8321_8323} -> *[]update_elem_8321_8321_8323" ~ "a_8329_8327 -> b_8329_8328",+        "b_8329_8328" ~ "a_8329_8325",+        "{xs: []t_8321_8322_8324} -> [][]t_8321_8322_8324" ~ "a_8329_8325 -> b_8329_8326",+        "b_8329_8326" ~ "a_8329_8323",+        "{a: [][]t_8321_8322_8329} -> [][]t_8321_8322_8329" ~ "a_8329_8323 -> b_8329_8324",+        "t_8321_8323_8324" ~ "b_8329_8324",+        "i64" ~ "t_8321_8323_8327",+        "i64" ~ "t_8321_8323_8327",+        "t_8321_8323_8327" ~ "t_8321_8323_8326",+        "i64" ~ "t_8321_8323_8326",+        "t_8321_8324_8327" ~ "t_8321_8324_8326",+        "index_8321_8324_8328" ~ "index_elem_8321_8324_8329",+        "t_8328_8325" ~ "[]index_elem_8321_8324_8329",+        "t_8321_8325_8320" ~ "index_8321_8324_8328",+        "num_8321_8325_8322" ~ "i32",+        "i32 -> t_1 -> i32" ~ "i32 -> t_8321_8325_8321 -> i32",+        "i32" ~ "i32",+        "t_8321_8325_8320" ~ "t_8321_8325_8321",+        "t_8321_8326_8321" ~ "i64",+        "t_8321_8324_8326" ~ "t_8321_8326_8322",+        "i64" ~ "t_8321_8326_8322",+        "t_8321_8326_8327" ~ "t_8321_8326_8322",+        "t_8321_8326_8328" ~ "t_8321_8326_8327",+        "index_8321_8326_8329" ~ "index_elem_8321_8327_8320",+        "t_8321_8323_8324" ~ "[]index_elem_8321_8327_8320",+        "t_8321_8327_8321" ~ "t_8321_8326_8321",+        "index_8321_8327_8322" ~ "index_elem_8321_8327_8323",+        "index_8321_8326_8329" ~ "[]index_elem_8321_8327_8323",+        "t_8321_8327_8324" ~ "index_8321_8327_8322",+        "t_8321_8327_8325" ~ "t_8321_8326_8327",+        "index_8321_8327_8326" ~ "index_elem_8321_8327_8327",+        "t_8329_8322" ~ "[]index_elem_8321_8327_8327",+        "t_8321_8327_8328" ~ "t_8321_8326_8321",+        "index_8321_8327_8329" ~ "index_elem_8321_8328_8320",+        "index_8321_8327_8326" ~ "[]index_elem_8321_8328_8320",+        "index_8321_8327_8329" ~ "i16",+        "t_8321_8328_8323" ~ "i64",+        "i64" ~ "t_8321_8328_8325",+        "t_8321_8326_8327" ~ "t_8321_8328_8325",+        "t_8321_8328_8325" ~ "t_8321_8328_8324",+        "t_8321_8328_8323" ~ "t_8321_8328_8324",+        "t_8321_8329_8324" ~ "t_8321_8328_8324",+        "t_8321_8324_8326" ~ "t_8321_8329_8326",+        "t_8321_8329_8324" ~ "t_8321_8329_8326",+        "t_8321_8329_8326" ~ "t_8321_8329_8325",+        "t_8321_8327_8324" ~ "t_8321_8329_8325",+        "t_8322_8320_8325" ~ "t_8321_8329_8325",+        "t_8321_8323_8326" ~ "i64",+        "{i: t_8321_8324_8326} -> t_8322_8320_8325" ~ "i64 -> a_8321_8323_8325",+        "t_8322_8321_8320" ~ "[]a_8321_8323_8325",+        "[]t_1" ~ "[]t_8322_8321_8321",+        "t_8322_8321_8320" ~ "[]i64",+        "t_8328_8325" ~ "[]t_8322_8321_8321"+      ],+      M.fromList [("t_1", (0, Unlifted, NoLoc))],+      M.fromList [("t_8320", (3, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321", (3, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8320", (4, TyVarFree NoLoc Unlifted)), ("t_8321_8323", (5, TyVarFree NoLoc Lifted)), ("t_8321_8324", (5, TyVarFree NoLoc Lifted)), ("t_8321_8325", (6, TyVarFree NoLoc Unlifted)), ("t_8322_8322", (7, TyVarFree NoLoc Lifted)), ("t_8322_8323", (7, TyVarFree NoLoc Lifted)), ("t_8322_8324", (7, TyVarFree NoLoc Lifted)), ("a_8322_8325", (8, TyVarFree NoLoc Lifted)), ("b_8322_8326", (8, TyVarFree NoLoc Lifted)), ("a_8322_8327", (8, TyVarFree NoLoc Lifted)), ("b_8322_8328", (8, TyVarFree NoLoc Lifted)), ("t_8322_8329", (8, TyVarFree NoLoc Unlifted)), ("a_8323_8322", (8, TyVarFree NoLoc Unlifted)), ("x_8323_8323", (8, TyVarFree NoLoc Unlifted)), ("t_8323_8324", (8, TyVarFree NoLoc Unlifted)), ("a_8324_8325", (8, TyVarFree NoLoc Unlifted)), ("b_8324_8326", (8, TyVarFree NoLoc Unlifted)), ("c_8324_8327", (8, TyVarFree NoLoc Unlifted)), ("t_8325_8322", (9, TyVarFree NoLoc Lifted)), ("t_8325_8323", (9, TyVarFree NoLoc Lifted)), ("t_8325_8324", (9, TyVarFree NoLoc Lifted)), ("t_8325_8325", (10, TyVarFree NoLoc Unlifted)), ("et_8325_8326", (10, TyVarFree NoLoc Unlifted)), ("t_8326_8321", (11, TyVarFree NoLoc Lifted)), ("t_8326_8322", (12, TyVarFree NoLoc Unlifted)), ("t_8326_8323", (12, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8326_8324", (12, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8327_8323", (12, TyVarFree NoLoc Unlifted)), ("t_8327_8324", (12, TyVarFree NoLoc Unlifted)), ("t_8328_8325", (13, TyVarFree NoLoc Lifted)), ("t_8328_8326", (14, TyVarFree NoLoc Unlifted)), ("et_8328_8327", (14, TyVarFree NoLoc Unlifted)), ("t_8329_8322", (15, TyVarFree NoLoc Lifted)), ("a_8329_8323", (16, TyVarFree NoLoc Lifted)), ("b_8329_8324", (16, TyVarFree NoLoc Lifted)), ("a_8329_8325", (16, TyVarFree NoLoc Lifted)), ("b_8329_8326", (16, TyVarFree NoLoc Lifted)), ("a_8329_8327", (16, TyVarFree NoLoc Lifted)), ("b_8329_8328", (16, TyVarFree NoLoc Lifted)), ("a_8329_8329", (16, TyVarFree NoLoc Lifted)), ("b_8321_8320_8320", (16, TyVarFree NoLoc Lifted)), ("a_8321_8320_8321", (16, TyVarFree NoLoc Lifted)), ("b_8321_8320_8322", (16, TyVarFree NoLoc Lifted)), ("t_8321_8320_8323", (16, TyVarFree NoLoc Unlifted)), ("t_8321_8320_8328", (16, TyVarFree NoLoc Unlifted)), ("update_elem_8321_8321_8323", (16, TyVarFree NoLoc Unlifted)), ("t_8321_8321_8324", (16, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8321_8325", (16, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8322_8324", (16, TyVarFree NoLoc Unlifted)), ("t_8321_8322_8329", (16, TyVarFree NoLoc Unlifted)), ("t_8321_8323_8324", (17, TyVarFree NoLoc Lifted)), ("a_8321_8323_8325", (18, TyVarFree NoLoc Unlifted)), ("t_8321_8323_8326", (18, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8323_8327", (18, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8324_8326", (19, TyVarFree NoLoc Lifted)), ("t_8321_8324_8327", (20, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8321_8324_8328", (20, TyVarFree NoLoc Unlifted)), ("index_elem_8321_8324_8329", (20, TyVarFree NoLoc Unlifted)), ("t_8321_8325_8320", (21, TyVarFree NoLoc Lifted)), ("t_8321_8325_8321", (22, TyVarFree NoLoc Unlifted)), ("num_8321_8325_8322", (22, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8326_8321", (23, TyVarFree NoLoc Lifted)), ("t_8321_8326_8322", (24, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8326_8327", (25, TyVarFree NoLoc Lifted)), ("t_8321_8326_8328", (26, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8321_8326_8329", (26, TyVarFree NoLoc Unlifted)), ("index_elem_8321_8327_8320", (26, TyVarFree NoLoc Unlifted)), ("t_8321_8327_8321", (26, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8321_8327_8322", (26, TyVarFree NoLoc Unlifted)), ("index_elem_8321_8327_8323", (26, TyVarFree NoLoc Unlifted)), ("t_8321_8327_8324", (27, TyVarFree NoLoc Lifted)), ("t_8321_8327_8325", (28, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8321_8327_8326", (28, TyVarFree NoLoc Unlifted)), ("index_elem_8321_8327_8327", (28, TyVarFree NoLoc Unlifted)), ("t_8321_8327_8328", (28, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8321_8327_8329", (28, TyVarFree NoLoc Unlifted)), ("index_elem_8321_8328_8320", (28, TyVarFree NoLoc Unlifted)), ("t_8321_8328_8323", (29, TyVarFree NoLoc Lifted)), ("t_8321_8328_8324", (30, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8328_8325", (30, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8329_8324", (31, TyVarFree NoLoc Lifted)), ("t_8321_8329_8325", (32, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8329_8326", (32, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8320_8325", (33, TyVarFree NoLoc Lifted)), ("t_8322_8321_8320", (19, TyVarFree NoLoc Lifted)), ("t_8322_8321_8321", (20, TyVarFree NoLoc Unlifted)), ("t_8328_8325_8322_8321_8328", (13, TyVarFree NoLoc Lifted)), ("t_8322_8321_8320_8322_8321_8329", (19, TyVarFree NoLoc Lifted)), ("index_8321_8327_8326_8322_8322_8320", (28, TyVarFree NoLoc Unlifted)), ("t_8329_8322_8322_8322_8321", (15, TyVarFree NoLoc Lifted)), ("index_elem_8321_8327_8327_8322_8322_8322", (28, TyVarFree NoLoc Unlifted)), ("index_8321_8326_8329_8322_8322_8323", (26, TyVarFree NoLoc Unlifted)), ("t_8321_8323_8324_8322_8322_8324", (17, TyVarFree NoLoc Lifted)), ("index_elem_8321_8327_8320_8322_8322_8325", (26, TyVarFree NoLoc Unlifted)), ("b_8329_8324_8322_8322_8326", (16, TyVarFree NoLoc Lifted)), ("a_8329_8323_8322_8322_8327", (16, TyVarFree NoLoc Lifted)), ("b_8329_8326_8322_8322_8328", (16, TyVarFree NoLoc Lifted)), ("a_8329_8325_8322_8322_8329", (16, TyVarFree NoLoc Lifted)), ("b_8329_8328_8322_8323_8320", (16, TyVarFree NoLoc Lifted)), ("a_8329_8327_8322_8323_8321", (16, TyVarFree NoLoc Lifted)), ("b_8321_8320_8320_8322_8323_8322", (16, TyVarFree NoLoc Lifted)), ("a_8329_8329_8322_8323_8323", (16, TyVarFree NoLoc Lifted)), ("b_8321_8320_8322_8322_8323_8324", (16, TyVarFree NoLoc Lifted)), ("a_8321_8320_8321_8322_8323_8325", (16, TyVarFree NoLoc Lifted)), ("t_8326_8321_8322_8323_8326", (11, TyVarFree NoLoc Lifted)), ("t_8328_8326_8322_8323_8327", (14, TyVarFree NoLoc Unlifted)), ("et_8328_8327_8322_8323_8328", (14, TyVarFree NoLoc Unlifted)), ("t_8325_8324_8322_8323_8329", (9, TyVarFree NoLoc Lifted)), ("t_8322_8324_8322_8324_8320", (7, TyVarFree NoLoc Lifted)), ("t_8322_8322_8322_8324_8321", (7, TyVarFree NoLoc Lifted)), ("t_8325_8322_8322_8324_8322", (9, TyVarFree NoLoc Lifted)), ("t_8325_8325_8322_8324_8323", (10, TyVarFree NoLoc Unlifted)), ("et_8325_8326_8322_8324_8324", (10, TyVarFree NoLoc Unlifted)), ("t_8325_8323_8322_8324_8325", (9, TyVarFree NoLoc Lifted)), ("t_8322_8323_8322_8324_8326", (7, TyVarFree NoLoc Lifted)), ("a_8322_8325_8322_8324_8327", (8, TyVarFree NoLoc Lifted)), ("b_8322_8328_8322_8324_8328", (8, TyVarFree NoLoc Lifted)), ("a_8323_8322_8322_8324_8329", (8, TyVarFree NoLoc Unlifted)), ("a_8322_8327_8322_8325_8320", (8, TyVarFree NoLoc Lifted)), ("t_8321_8323_8322_8325_8321", (5, TyVarFree NoLoc Lifted)), ("t_8321_8324_8322_8325_8322", (5, TyVarFree NoLoc Lifted))]+    ),+    ( [ "b_8325_8329" ~ "[]b_8325_8329_8328_8325",+        "t_8326_8326" ~ "[]t_8326_8326_8328_8326",+        "a_8325_8328" ~ "[]a_8325_8328_8328_8327",+        "t_8326_8323" ~ "[]t_8326_8323_8328_8328",+        "t_8325_8327" ~ "[]t_8325_8327_8328_8329",+        "i64" ~ "t_8320",+        "num_8321" ~ "t_8320",+        "i32" ~ "t_8329",+        "num_8321_8320" ~ "t_8329",+        "t_8329" ~ "t_8328",+        "num_8321_8325" ~ "t_8328",+        "t_8328" ~ "t_8327",+        "num_8322_8320" ~ "t_8327",+        "bool" ~ "bool",+        "num_8326" ~ "if_t_8322_8325",+        "t_8327" ~ "if_t_8322_8325",+        "t_8322_8326" ~ "if_t_8322_8325",+        "i16" ~ "i16",+        "t_8322_8329" ~ "i64",+        "i64" ~ "t_8323_8320",+        "t_8322_8329" ~ "t_8323_8320",+        "t_8323_8325" ~ "t_8323_8320",+        "i64" ~ "t_8323_8326",+        "t_8322_8329" ~ "t_8323_8326",+        "t_8324_8321" ~ "t_8323_8326",+        "t_8323_8325" ~ "t_8324_8324",+        "t_8322_8329" ~ "t_8324_8324",+        "t_8324_8324" ~ "t_8324_8323",+        "t_8324_8321" ~ "t_8324_8323",+        "t_8324_8323" ~ "i64",+        "[]t_1" ~ "[]t_8324_8322",+        "t_8325_8327" ~ "[]t_8324_8322",+        "i64" ~ "i64",+        "t_8325_8327" ~ "t_8326_8323",+        "t_8326_8326" ~ "t_8326_8323",+        "t_8326_8327" ~ "t_8322_8326",+        "t_8322_8326" ~ "t_8326_8329",+        "num_8327_8320" ~ "t_8326_8329",+        "i32 -> t_1 -> i32" ~ "i32 -> t_8326_8328 -> i32",+        "t_8326_8329" ~ "i32",+        "t_8326_8326" ~ "[]t_8326_8328",+        "t_8326_8323" ~ "[]t_8326_8328",+        "{xs: []t_8326_8320} -> []t_8326_8320" ~ "a_8325_8328 -> b_8325_8329",+        "t_8326_8326" ~ "a_8325_8328",+        "[]t_1" ~ "b_8325_8329"+      ],+      M.fromList [("t_1", (0, Unlifted, NoLoc))],+      M.fromList [("t_8320", (5, TyVarFree NoLoc Unlifted)), ("num_8321", (5, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8326", (5, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8327", (5, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8328", (5, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8329", (5, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8320", (5, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8325", (5, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8320", (5, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("if_t_8322_8325", (5, TyVarFree NoLoc SizeLifted)), ("t_8322_8326", (6, TyVarFree NoLoc Lifted)), ("t_8322_8329", (8, TyVarFree NoLoc Lifted)), ("t_8323_8320", (9, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8325", (10, TyVarFree NoLoc Lifted)), ("t_8323_8326", (11, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8324_8321", (12, TyVarFree NoLoc Lifted)), ("t_8324_8322", (13, TyVarFree NoLoc Unlifted)), ("t_8324_8323", (13, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8324_8324", (13, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8325_8327", (14, TyVarFree NoLoc Lifted)), ("a_8325_8328", (15, TyVarFree NoLoc Lifted)), ("b_8325_8329", (15, TyVarFree NoLoc Lifted)), ("t_8326_8320", (15, TyVarFree NoLoc Unlifted)), ("t_8326_8323", (15, TyVarFree NoLoc Unlifted)), ("t_8326_8326", (15, TyVarFree NoLoc Lifted)), ("t_8326_8327", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64])), ("t_8326_8328", (15, TyVarFree NoLoc Unlifted)), ("t_8326_8329", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8327_8320", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("b_8325_8329_8328_8325", (15, TyVarFree NoLoc Lifted)), ("t_8326_8326_8328_8326", (15, TyVarFree NoLoc Lifted)), ("a_8325_8328_8328_8327", (15, TyVarFree NoLoc Lifted)), ("t_8326_8323_8328_8328", (15, TyVarFree NoLoc Unlifted)), ("t_8325_8327_8328_8329", (14, TyVarFree NoLoc Lifted))]+    ),+    ( [ "t_8324_8325_8320" ~ "[]t_8324_8325_8320_8324_8326_8326",+        "t_8323_8327_8320" ~ "[]t_8323_8327_8320_8324_8326_8327",+        "t_8323_8327_8322" ~ "[]t_8323_8327_8322_8324_8326_8328",+        "t_8323_8327_8321" ~ "[]t_8323_8327_8321_8324_8326_8329",+        "a_8323_8325_8320" ~ "[]a_8323_8325_8320_8324_8327_8320",+        "t_8321" ~ "t_8323",+        "num_8324" ~ "t_8323",+        "i64" ~ "i64",+        "t_8321" ~ "t_8329",+        "i32" ~ "t_8329",+        "bool" ~ "t_8322",+        "bool" ~ "t_8322",+        "t_8322_8320" ~ "t_8320",+        "index_8322_8321" ~ "index_elem_8322_8322",+        "[]u32" ~ "[]index_elem_8322_8322",+        "t_8322_8323" ~ "index_8322_8321",+        "t_8322_8324" ~ "t_8321",+        "index_8322_8325" ~ "index_elem_8322_8326",+        "[]u32" ~ "[]index_elem_8322_8326",+        "t_8322_8327" ~ "index_8322_8325",+        "t_8322_8323" ~ "t_8322_8328",+        "t_8322_8327" ~ "t_8322_8328",+        "t_8320" ~ "i32",+        "t_8321" ~ "i32",+        "u32" ~ "t_8323_8325",+        "u32" ~ "t_8323_8325",+        "t_8323_8325" ~ "u32",+        "num_8323_8324" ~ "t_8323_8323",+        "i32" ~ "t_8323_8323",+        "t_8322_8323" ~ "t_8325_8320",+        "t_8322_8327" ~ "t_8325_8320",+        "t_8325_8320" ~ "u32",+        "bool" ~ "bool",+        "t_8323_8323" ~ "if_t_8325_8327",+        "i32" ~ "if_t_8325_8327",+        "t_8325_8329" ~ "num_8325_8328",+        "t_8322" ~ "bool",+        "if_t_8325_8327" ~ "if_t_8326_8320",+        "num_8325_8328" ~ "if_t_8326_8320",+        "t_8326_8321" ~ "i64",+        "t_8326_8324" ~ "i32",+        "t_8326_8324" ~ "t_8326_8326",+        "num_8326_8327" ~ "t_8326_8326",+        "(t_8326_8324, t_8326_8326)" ~ "(t_8320, t_8321)",+        "t_8326_8324" ~ "t_8327_8324",+        "num_8327_8325" ~ "t_8327_8324",+        "(t_8326_8324, t_8327_8324)" ~ "(t_8320, t_8321)",+        "if_t_8326_8320" ~ "t_8326_8325",+        "if_t_8326_8320" ~ "t_8326_8325",+        "t_8326_8325" ~ "i32",+        "t_8328_8328" ~ "i32",+        "t_8326_8324" ~ "t_8328_8329",+        "t_8328_8328" ~ "t_8328_8329",+        "(t_8326_8324, t_8328_8329)" ~ "(t_8320, t_8321)",+        "t_8329_8326" ~ "if_t_8326_8320",+        "t_8329_8328" ~ "num_8329_8327",+        "t_8329_8328" ~ "t_8321_8320_8321",+        "t_8328_8328" ~ "t_8321_8320_8321",+        "t_8326_8324" ~ "t_8321_8320_8320",+        "t_8321_8320_8321" ~ "t_8321_8320_8320",+        "(t_8326_8324, t_8321_8320_8320)" ~ "(t_8320, t_8321)",+        "if_t_8326_8320" ~ "t_8329_8329",+        "t_8329_8326" ~ "t_8329_8329",+        "t_8329_8328" ~ "t_8321_8321_8326",+        "num_8321_8321_8327" ~ "t_8321_8321_8326",+        "num_8329_8327" ~ "t_8321_8321_8326",+        "t_8321_8322_8322" ~ "t_8329_8328",+        "t_8321_8322_8322" ~ "t_8321_8322_8324",+        "num_8321_8322_8325" ~ "t_8321_8322_8324",+        "(t_8321_8323_8320, t_8321_8323_8321)" ~ "(num_8321_8322_8323, t_8321_8322_8324)",+        "t_8321_8323_8321" ~ "t_8321_8323_8322",+        "num_8321_8323_8323" ~ "t_8321_8323_8322",+        "t_8321_8323_8320" ~ "t_8321_8324_8321",+        "t_8321_8323_8321" ~ "t_8321_8324_8321",+        "t_8321_8324_8321" ~ "t_8321_8324_8320",+        "t_8328_8328" ~ "t_8321_8324_8320",+        "t_8326_8324" ~ "t_8321_8323_8329",+        "t_8321_8324_8320" ~ "t_8321_8323_8329",+        "(t_8326_8324, t_8321_8323_8329)" ~ "(t_8320, t_8321)",+        "if_t_8326_8320" ~ "t_8321_8323_8328",+        "t_8329_8326" ~ "t_8321_8323_8328",+        "t_8321_8323_8320" ~ "t_8321_8326_8320",+        "t_8321_8323_8321" ~ "t_8321_8326_8320",+        "t_8321_8323_8321" ~ "t_8321_8326_8325",+        "num_8321_8326_8326" ~ "t_8321_8326_8325",+        "t_8321_8323_8321" ~ "t_8321_8327_8321",+        "num_8321_8327_8322" ~ "t_8321_8327_8321",+        "bool" ~ "bool",+        "(t_8321_8326_8320, t_8321_8326_8325)" ~ "if_t_8321_8327_8327",+        "(t_8321_8323_8320, t_8321_8327_8321)" ~ "if_t_8321_8327_8327",+        "(num_8321_8322_8323, t_8321_8322_8324)" ~ "if_t_8321_8327_8327",+        "(t_8321_8327_8328, t_8321_8327_8329)" ~ "(t_8321_8323_8320, t_8321_8323_8321)",+        "t_8321_8327_8328" ~ "t_8321_8328_8321",+        "t_8328_8328" ~ "t_8321_8328_8321",+        "t_8326_8324" ~ "t_8321_8328_8320",+        "t_8321_8328_8321" ~ "t_8321_8328_8320",+        "t_8321_8329_8320" ~ "t_8321_8328_8320",+        "(t_8326_8324, t_8321_8329_8320)" ~ "(t_8320, t_8321)",+        "t_8321_8329_8323" ~ "if_t_8326_8320",+        "(t_8321_8329_8326, t_8321_8329_8327)" ~ "(num_8321_8329_8324, num_8321_8329_8325)",+        "t_8321_8329_8327" ~ "t_8321_8329_8328",+        "t_8321_8327_8328" ~ "t_8321_8329_8328",+        "t_8321_8329_8327" ~ "t_8322_8320_8323",+        "num_8322_8320_8324" ~ "t_8322_8320_8323",+        "t_8321_8327_8328" ~ "i32",+        "t_8322_8320_8323" ~ "i32",+        "t_8322_8321_8323" ~ "i32",+        "t_8321_8329_8326" ~ "t_8322_8321_8327",+        "t_8322_8321_8323" ~ "t_8322_8321_8327",+        "t_8322_8321_8327" ~ "t_8322_8321_8326",+        "t_8328_8328" ~ "t_8322_8321_8326",+        "t_8326_8324" ~ "t_8322_8321_8325",+        "t_8322_8321_8326" ~ "t_8322_8321_8325",+        "(t_8326_8324, t_8322_8321_8325)" ~ "(t_8320, t_8321)",+        "if_t_8326_8320" ~ "t_8322_8321_8324",+        "t_8321_8329_8323" ~ "t_8322_8321_8324",+        "t_8321_8329_8326" ~ "t_8322_8323_8326",+        "t_8322_8321_8323" ~ "t_8322_8323_8326",+        "t_8321_8329_8327" ~ "t_8322_8324_8321",+        "num_8322_8324_8322" ~ "t_8322_8324_8321",+        "t_8321_8329_8327" ~ "t_8322_8324_8327",+        "num_8322_8324_8328" ~ "t_8322_8324_8327",+        "bool" ~ "bool",+        "(t_8322_8323_8326, t_8322_8324_8321)" ~ "if_t_8322_8325_8323",+        "(t_8321_8329_8326, t_8322_8324_8327)" ~ "if_t_8322_8325_8323",+        "(num_8321_8329_8324, num_8321_8329_8325)" ~ "if_t_8322_8325_8323",+        "(t_8322_8325_8324, t_8322_8325_8325)" ~ "(t_8321_8329_8326, t_8321_8329_8327)",+        "t_8322_8325_8324" ~ "t_8322_8325_8328",+        "t_8328_8328" ~ "t_8322_8325_8328",+        "t_8326_8324" ~ "t_8322_8325_8327",+        "t_8322_8325_8328" ~ "t_8322_8325_8327",+        "t_8328_8328" ~ "i32",+        "num_8322_8326_8327" ~ "i32",+        "t_8322_8325_8327" ~ "t_8322_8325_8326",+        "i32" ~ "t_8322_8325_8326",+        "t_8322_8327_8326" ~ "t_8322_8325_8326",+        "t_8326_8324" ~ "i32",+        "t_8321_8329_8320" ~ "i32",+        "i32" ~ "t_8322_8327_8327",+        "t_8322_8327_8326" ~ "t_8322_8327_8327",+        "t_8322_8328_8328" ~ "num_8322_8328_8327",+        "bool" ~ "bool",+        "(t_8322_8328_8326, num_8322_8328_8327)" ~ "if_t_8322_8329_8320",+        "(t_8322_8328_8329, t_8322_8327_8326)" ~ "if_t_8322_8329_8320",+        "(t_8322_8329_8321, t_8322_8329_8322)" ~ "if_t_8322_8329_8320",+        "t_8326_8324" ~ "i32",+        "t_8321_8329_8320" ~ "i32",+        "t_8322_8327_8326" ~ "t_8322_8329_8328",+        "num_8322_8329_8329" ~ "t_8322_8329_8328",+        "i32" ~ "t_8322_8329_8323",+        "t_8322_8329_8328" ~ "t_8322_8329_8323",+        "t_8322_8327_8326" ~ "t_8323_8320_8328",+        "num_8323_8320_8329" ~ "t_8323_8320_8328",+        "t_8323_8321_8326" ~ "num_8323_8321_8325",+        "t_8322_8327_8326" ~ "t_8323_8321_8327",+        "num_8323_8321_8328" ~ "t_8323_8321_8327",+        "t_8322_8327_8326" ~ "t_8323_8322_8324",+        "num_8323_8322_8325" ~ "t_8323_8322_8324",+        "bool" ~ "bool",+        "(t_8323_8321_8324, num_8323_8321_8325)" ~ "if_t_8323_8323_8320",+        "(t_8323_8322_8323, t_8323_8322_8324)" ~ "if_t_8323_8323_8320",+        "(t_8323_8323_8321, t_8323_8323_8322)" ~ "if_t_8323_8323_8320",+        "t_8323_8323_8324" ~ "t_8326_8324",+        "index_8323_8323_8325" ~ "index_elem_8323_8323_8326",+        "[]u32" ~ "[]index_elem_8323_8323_8326",+        "t_8321_8329_8323" ~ "i32",+        "num_8323_8323_8328" ~ "t_8323_8323_8327",+        "u32" ~ "t_8323_8323_8327",+        "index_8323_8323_8325" ~ "t_8323_8323_8323",+        "t_8323_8323_8327" ~ "t_8323_8323_8323",+        "t_8323_8324_8329" ~ "t_8323_8323_8323",+        "i64" ~ "t_8323_8325_8323",+        "num_8323_8325_8324" ~ "t_8323_8325_8323",+        "t_8323_8325_8323" ~ "i64",+        "{i: t_8326_8321} -> ({delta_node: t_8321_8329_8323, left: t_8322_8329_8321, right: t_8323_8323_8321, sfc_code: t_8323_8324_8329}, (t_8322_8329_8322, t_8326_8324), (t_8323_8323_8322, t_8326_8324))" ~ "i64 -> a_8323_8325_8322",+        "[]a_8323_8325_8322" ~ "a_8323_8325_8320",+        "{xs: [](a_8323_8326_8323, b_8323_8326_8324, c_8323_8326_8325)} -> ([]a_8323_8326_8323, []b_8323_8326_8324, []c_8323_8326_8325)" ~ "a_8323_8325_8320 -> b_8323_8325_8321",+        "(t_8323_8327_8320, t_8323_8327_8321, t_8323_8327_8322)" ~ "b_8323_8325_8321",+        "i64" ~ "t_8323_8327_8324",+        "num_8323_8327_8325" ~ "t_8323_8327_8324",+        "t_8323_8328_8321" ~ "num_8323_8328_8320",+        "a_8323_8328_8328" ~ "ft_8323_8329_8320",+        "a_8323_8328_8328 -> b_8323_8328_8329" ~ "a_8323_8328_8325 -> b_8323_8328_8326",+        "i32 -> i64" ~ "b_8323_8328_8326 -> c_8323_8328_8327",+        "{x: a_8323_8328_8325} -> c_8323_8328_8327" ~ "a_8323_8328_8323 -> x_8323_8328_8324",+        "t_8323_8327_8321" ~ "[]a_8323_8328_8323",+        "a_8324_8320_8324" ~ "ft_8324_8320_8326",+        "a_8324_8320_8324 -> b_8324_8320_8325" ~ "a_8324_8320_8321 -> b_8324_8320_8322",+        "i32 -> i64" ~ "b_8324_8320_8322 -> c_8324_8320_8323",+        "{x: a_8324_8320_8321} -> c_8324_8320_8323" ~ "a_8323_8329_8329 -> x_8324_8320_8320",+        "t_8323_8327_8322" ~ "[]a_8323_8329_8329",+        "[]x_8323_8328_8324" ~ "[]t_8323_8328_8322",+        "[]x_8324_8320_8320" ~ "[]t_8323_8328_8322",+        "a_8324_8322_8322" ~ "ft_8324_8322_8324",+        "a_8324_8322_8322 -> b_8324_8322_8323" ~ "a_8324_8322_8320 -> x_8324_8322_8321",+        "t_8323_8327_8321" ~ "[]a_8324_8322_8320",+        "a_8324_8323_8321" ~ "ft_8324_8323_8323",+        "a_8324_8323_8321 -> b_8324_8323_8322" ~ "a_8324_8322_8329 -> x_8324_8323_8320",+        "t_8323_8327_8322" ~ "[]a_8324_8322_8329",+        "[]x_8324_8322_8321" ~ "[]t_8324_8321_8329",+        "[]x_8324_8323_8320" ~ "[]t_8324_8321_8329",+        "t_8323_8327_8324" ~ "i64",+        "num_8323_8328_8320" ~ "t_8323_8327_8323",+        "[]t_8323_8328_8322" ~ "[]i64",+        "[]t_8324_8321_8329" ~ "[]t_8323_8327_8323",+        "t_8324_8325_8320" ~ "[]t_8323_8327_8323",+        "{delta_node: t_8324_8325_8324, left: t_8324_8325_8325, right: t_8324_8325_8326, sfc_code: t_8324_8325_8327} -> {parent: t_8324_8325_8328} -> {delta_node: t_8324_8325_8324, left: t_8324_8325_8325, parent: t_8324_8325_8328, right: t_8324_8325_8326, sfc_code: t_8324_8325_8327}" ~ "a_8324_8325_8321 -> b_8324_8325_8322 -> x_8324_8325_8323",+        "t_8323_8327_8320" ~ "[]a_8324_8325_8321",+        "t_8324_8325_8320" ~ "[]b_8324_8325_8322",+        "[]{delta_node: i32, left: #inner i32 | #leaf i32, parent: i32, right: #inner i32 | #leaf i32, sfc_code: u32}" ~ "[]x_8324_8325_8323"+      ],+      M.empty,+      M.fromList [("t_8320", (3, TyVarFree NoLoc Lifted)), ("t_8321", (3, TyVarFree NoLoc Lifted)), ("t_8322", (4, TyVarPrim NoLoc [Bool])), ("t_8323", (4, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8324", (4, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8329", (4, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8320", (4, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8322_8321", (4, TyVarFree NoLoc Unlifted)), ("index_elem_8322_8322", (4, TyVarFree NoLoc Unlifted)), ("t_8322_8323", (5, TyVarFree NoLoc Lifted)), ("t_8322_8324", (6, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8322_8325", (6, TyVarFree NoLoc Unlifted)), ("index_elem_8322_8326", (6, TyVarFree NoLoc Unlifted)), ("t_8322_8327", (7, TyVarFree NoLoc Lifted)), ("t_8322_8328", (8, TyVarFree NoLoc Unlifted)), ("t_8323_8323", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8324", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8325", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64])), ("t_8325_8320", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64])), ("if_t_8325_8327", (8, TyVarFree NoLoc SizeLifted)), ("num_8325_8328", (4, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8325_8329", (4, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("if_t_8326_8320", (4, TyVarFree NoLoc SizeLifted)), ("t_8326_8321", (3, TyVarFree NoLoc Lifted)), ("t_8326_8324", (5, TyVarFree NoLoc Lifted)), ("t_8326_8325", (6, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8326_8326", (6, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8326_8327", (6, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8327_8324", (6, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8327_8325", (6, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8328_8328", (7, TyVarFree NoLoc Lifted)), ("t_8328_8329", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8329_8326", (9, TyVarFree NoLoc Lifted)), ("num_8329_8327", (10, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8329_8328", (10, TyVarFree NoLoc Lifted)), ("t_8329_8329", (10, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8320_8320", (10, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8320_8321", (10, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8321_8326", (10, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8321_8327", (10, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8322_8322", (11, TyVarFree NoLoc Lifted)), ("num_8321_8322_8323", (12, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8322_8324", (12, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8322_8325", (12, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8323_8320", (12, TyVarFree NoLoc Lifted)), ("t_8321_8323_8321", (12, TyVarFree NoLoc Lifted)), ("t_8321_8323_8322", (12, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8323_8323", (12, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8323_8328", (12, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8323_8329", (12, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8324_8320", (12, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8324_8321", (12, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8326_8320", (12, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8326_8325", (12, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8326_8326", (12, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8327_8321", (12, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8327_8322", (12, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("if_t_8321_8327_8327", (12, TyVarFree NoLoc SizeLifted)), ("t_8321_8327_8328", (13, TyVarFree NoLoc Lifted)), ("t_8321_8327_8329", (13, TyVarFree NoLoc Lifted)), ("t_8321_8328_8320", (14, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8328_8321", (14, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8329_8320", (15, TyVarFree NoLoc Lifted)), ("t_8321_8329_8323", (17, TyVarFree NoLoc Lifted)), ("num_8321_8329_8324", (18, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8329_8325", (18, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8329_8326", (18, TyVarFree NoLoc Lifted)), ("t_8321_8329_8327", (18, TyVarFree NoLoc Lifted)), ("t_8321_8329_8328", (18, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8320_8323", (18, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8320_8324", (18, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8321_8323", (19, TyVarFree NoLoc Lifted)), ("t_8322_8321_8324", (20, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8321_8325", (20, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8321_8326", (20, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8321_8327", (20, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8323_8326", (20, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8324_8321", (20, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8324_8322", (20, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8324_8327", (20, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8324_8328", (20, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("if_t_8322_8325_8323", (20, TyVarFree NoLoc SizeLifted)), ("t_8322_8325_8324", (19, TyVarFree NoLoc Lifted)), ("t_8322_8325_8325", (19, TyVarFree NoLoc Lifted)), ("t_8322_8325_8326", (20, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8325_8327", (20, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8325_8328", (20, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8326_8327", (20, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8327_8326", (21, TyVarFree NoLoc Lifted)), ("t_8322_8327_8327", (22, TyVarFree NoLoc Unlifted)), ("t_8322_8328_8326", (22, TyVarSum NoLoc (M.fromList [("leaf", [Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "t_8322_8327_8326" 14782}) [])])]))), ("num_8322_8328_8327", (22, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8328_8328", (22, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8328_8329", (22, TyVarSum NoLoc (M.fromList [("inner", [Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "t_8322_8327_8326" 14782}) [])])]))), ("if_t_8322_8329_8320", (22, TyVarFree NoLoc SizeLifted)), ("t_8322_8329_8321", (23, TyVarFree NoLoc Lifted)), ("t_8322_8329_8322", (23, TyVarFree NoLoc Lifted)), ("t_8322_8329_8323", (24, TyVarFree NoLoc Unlifted)), ("t_8322_8329_8328", (24, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8329_8329", (24, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8320_8328", (24, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8320_8329", (24, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8321_8324", (24, TyVarSum NoLoc (M.fromList [("leaf", [Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "t_8323_8320_8328" 14819}) [])])]))), ("num_8323_8321_8325", (24, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8321_8326", (24, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8321_8327", (24, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8321_8328", (24, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8322_8323", (24, TyVarSum NoLoc (M.fromList [("inner", [Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "t_8323_8321_8327" 14829}) [])])]))), ("t_8323_8322_8324", (24, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8322_8325", (24, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("if_t_8323_8323_8320", (24, TyVarFree NoLoc SizeLifted)), ("t_8323_8323_8321", (25, TyVarFree NoLoc Lifted)), ("t_8323_8323_8322", (25, TyVarFree NoLoc Lifted)), ("t_8323_8323_8323", (26, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64])), ("t_8323_8323_8324", (26, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8323_8323_8325", (26, TyVarFree NoLoc Unlifted)), ("index_elem_8323_8323_8326", (26, TyVarFree NoLoc Unlifted)), ("t_8323_8323_8327", (26, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8323_8328", (26, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8324_8329", (27, TyVarFree NoLoc Lifted)), ("a_8323_8325_8320", (2, TyVarFree NoLoc Lifted)), ("b_8323_8325_8321", (2, TyVarFree NoLoc Lifted)), ("a_8323_8325_8322", (2, TyVarFree NoLoc Unlifted)), ("t_8323_8325_8323", (2, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8325_8324", (2, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("a_8323_8326_8323", (2, TyVarFree NoLoc Unlifted)), ("b_8323_8326_8324", (2, TyVarFree NoLoc Unlifted)), ("c_8323_8326_8325", (2, TyVarFree NoLoc Unlifted)), ("t_8323_8327_8320", (3, TyVarFree NoLoc Lifted)), ("t_8323_8327_8321", (3, TyVarFree NoLoc Lifted)), ("t_8323_8327_8322", (3, TyVarFree NoLoc Lifted)), ("t_8323_8327_8323", (4, TyVarFree NoLoc Unlifted)), ("t_8323_8327_8324", (4, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8327_8325", (4, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8328_8320", (4, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8328_8321", (4, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8328_8322", (4, TyVarFree NoLoc Unlifted)), ("a_8323_8328_8323", (4, TyVarFree NoLoc Unlifted)), ("x_8323_8328_8324", (4, TyVarFree NoLoc Unlifted)), ("a_8323_8328_8325", (4, TyVarFree NoLoc Lifted)), ("b_8323_8328_8326", (4, TyVarFree NoLoc Lifted)), ("c_8323_8328_8327", (4, TyVarFree NoLoc Lifted)), ("a_8323_8328_8328", (4, TyVarFree NoLoc Lifted)), ("b_8323_8328_8329", (4, TyVarFree NoLoc Lifted)), ("ft_8323_8329_8320", (4, TyVarRecord NoLoc (M.fromList [("0", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "b_8323_8328_8329" 14921}) []))]))), ("a_8323_8329_8329", (4, TyVarFree NoLoc Unlifted)), ("x_8324_8320_8320", (4, TyVarFree NoLoc Unlifted)), ("a_8324_8320_8321", (4, TyVarFree NoLoc Lifted)), ("b_8324_8320_8322", (4, TyVarFree NoLoc Lifted)), ("c_8324_8320_8323", (4, TyVarFree NoLoc Lifted)), ("a_8324_8320_8324", (4, TyVarFree NoLoc Lifted)), ("b_8324_8320_8325", (4, TyVarFree NoLoc Lifted)), ("ft_8324_8320_8326", (4, TyVarRecord NoLoc (M.fromList [("0", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "b_8324_8320_8325" 14941}) []))]))), ("t_8324_8321_8329", (4, TyVarFree NoLoc Unlifted)), ("a_8324_8322_8320", (4, TyVarFree NoLoc Unlifted)), ("x_8324_8322_8321", (4, TyVarFree NoLoc Unlifted)), ("a_8324_8322_8322", (4, TyVarFree NoLoc Lifted)), ("b_8324_8322_8323", (4, TyVarFree NoLoc Lifted)), ("ft_8324_8322_8324", (4, TyVarRecord NoLoc (M.fromList [("1", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "b_8324_8322_8323" 14967}) []))]))), ("a_8324_8322_8329", (4, TyVarFree NoLoc Unlifted)), ("x_8324_8323_8320", (4, TyVarFree NoLoc Unlifted)), ("a_8324_8323_8321", (4, TyVarFree NoLoc Lifted)), ("b_8324_8323_8322", (4, TyVarFree NoLoc Lifted)), ("ft_8324_8323_8323", (4, TyVarRecord NoLoc (M.fromList [("1", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "b_8324_8323_8322" 14979}) []))]))), ("t_8324_8325_8320", (5, TyVarFree NoLoc Lifted)), ("a_8324_8325_8321", (6, TyVarFree NoLoc Unlifted)), ("b_8324_8325_8322", (6, TyVarFree NoLoc Unlifted)), ("x_8324_8325_8323", (6, TyVarFree NoLoc Unlifted)), ("t_8324_8325_8324", (7, TyVarFree NoLoc Lifted)), ("t_8324_8325_8325", (7, TyVarFree NoLoc Lifted)), ("t_8324_8325_8326", (7, TyVarFree NoLoc Lifted)), ("t_8324_8325_8327", (7, TyVarFree NoLoc Lifted)), ("t_8324_8325_8328", (8, TyVarFree NoLoc Lifted)), ("t_8324_8325_8320_8324_8326_8326", (5, TyVarFree NoLoc Lifted)), ("t_8323_8327_8320_8324_8326_8327", (3, TyVarFree NoLoc Lifted)), ("t_8323_8327_8322_8324_8326_8328", (3, TyVarFree NoLoc Lifted)), ("t_8323_8327_8321_8324_8326_8329", (3, TyVarFree NoLoc Lifted)), ("a_8323_8325_8320_8324_8327_8320", (2, TyVarFree NoLoc Lifted))]+    ),+    ( [ "t_8326_8325_8325" ~ "[]t_8326_8325_8325_8326_8325_8327",+        "t_8322_8320_8326" ~ "[]t_8322_8320_8326_8326_8325_8328",+        "t_8326_8324_8326" ~ "[]t_8326_8324_8326_8326_8325_8329",+        "if_t_8326_8324_8325" ~ "[]if_t_8326_8324_8325_8326_8326_8320",+        "t_8324_8324_8326" ~ "[]t_8324_8324_8326_8326_8326_8321",+        "t_8326_8324_8320" ~ "[]t_8326_8324_8320_8326_8326_8322",+        "t_8324_8326_8325" ~ "[]t_8324_8326_8325_8326_8326_8323",+        "t_8324_8326_8326" ~ "[]t_8324_8326_8326_8326_8326_8324",+        "t_8326_8324_8321" ~ "[]t_8326_8324_8321_8326_8326_8325",+        "t_8324_8326_8323" ~ "[]t_8324_8326_8323_8326_8326_8326",+        "t_8327_8322" ~ "[]t_8327_8322_8326_8326_8327",+        "t_8321_8325_8326" ~ "[]t_8321_8325_8326_8326_8326_8328",+        "b_8321_8328_8321" ~ "[]b_8321_8328_8321_8326_8326_8329",+        "a_8321_8328_8320" ~ "[]a_8321_8328_8320_8326_8327_8320",+        "t_8321_8327_8329" ~ "[]t_8321_8327_8329_8326_8327_8321",+        "t_8321_8327_8328" ~ "[]t_8321_8327_8328_8326_8327_8322",+        "a_8321_8325_8327" ~ "[]a_8321_8325_8327_8326_8327_8323",+        "t_8321_8321_8329" ~ "[]t_8321_8321_8329_8326_8327_8324",+        "b_8327_8324" ~ "[]b_8327_8324_8326_8327_8325",+        "a_8327_8323" ~ "[]a_8327_8323_8326_8327_8326",+        "b_8327_8326" ~ "[]b_8327_8326_8326_8327_8327",+        "a_8327_8325" ~ "[]a_8327_8325_8326_8327_8328",+        "t_8321" ~ "num_8320",+        "index_8322" ~ "index_elem_8323",+        "[]{delta_node: i32, left: #inner i32 | #leaf i32, parent: i32, right: #inner i32 | #leaf i32, sfc_code: u32}" ~ "[]index_elem_8323",+        "index_8322" ~ "t_8325",+        "t_8326" ~ "kt_8324",+        "t_8327" ~ "t_8321_8320",+        "num_8321_8321" ~ "t_8321_8320",+        "t_8328" ~ "t_8321_8326",+        "num_8321_8327" ~ "t_8321_8326",+        "t_8321_8320" ~ "t_8329",+        "t_8321_8326" ~ "t_8329",+        "i32" ~ "t_8329",+        "t_8322_8328" ~ "t_8322_8326",+        "i32" ~ "t_8322_8329",+        "num_8323_8320" ~ "t_8322_8329",+        "t_8323_8326" ~ "t_8322_8326",+        "t_8323_8327" ~ "t_8323_8325",+        "index_8323_8328" ~ "index_elem_8323_8329",+        "[]{delta_node: i32, left: #inner i32 | #leaf i32, parent: i32, right: #inner i32 | #leaf i32, sfc_code: u32}" ~ "[]index_elem_8323_8329",+        "index_8323_8328" ~ "t_8324_8321",+        "t_8322_8329" ~ "kt_8324_8320",+        "match_t_8324_8322" ~ "t_8322_8329",+        "t_8324_8325" ~ "t_8324_8327",+        "kt_8324_8326" ~ "t_8322_8326",+        "t_8324_8325" ~ "t_8325_8321",+        "match_t_8324_8322" ~ "t_8327",+        "kt_8325_8320" ~ "t_8328",+        "t_8325_8326" ~ "i32",+        "t_8324_8325" ~ "t_8325_8328",+        "kt_8325_8327" ~ "t_8322_8326",+        "t_8324_8325" ~ "t_8326_8322",+        "match_t_8324_8322" ~ "t_8327",+        "kt_8326_8321" ~ "t_8328",+        "t_8326_8327" ~ "i32",+        "{n: t_8324_8325} -> {left: t_8325_8326, right: t_8326_8327}" ~ "a_8324_8323 -> x_8324_8324",+        "[]{delta_node: i32, left: #inner i32 | #leaf i32, parent: i32, right: #inner i32 | #leaf i32, sfc_code: u32}" ~ "[]a_8324_8323",+        "t_8327_8322" ~ "[]x_8324_8324",+        "t_8327_8329" ~ "t_8328_8322",+        "num_8328_8323" ~ "t_8328_8322",+        "bool" ~ "bool",+        "t_8328_8320" ~ "t_8329_8320",+        "num_8329_8321" ~ "t_8329_8320",+        "bool" ~ "bool",+        "i32" ~ "t_8328_8321",+        "i32" ~ "t_8328_8321",+        "{left: t_8327_8329, right: t_8328_8320} -> t_8328_8321" ~ "a_8327_8327 -> x_8327_8328",+        "t_8327_8322" ~ "a_8327_8325",+        "{as: []a_8327_8327} -> *[]x_8327_8328" ~ "a_8327_8325 -> b_8327_8326",+        "t_8321_8320_8329 -> t_8321_8320_8329 -> t_8321_8320_8329" ~ "a_8321_8320_8328 -> a_8321_8320_8328 -> a_8321_8320_8328",+        "num_8321_8321_8320" ~ "a_8321_8320_8328",+        "b_8327_8326" ~ "a_8327_8323",+        "{as: []a_8321_8320_8328} -> *[]a_8321_8320_8328" ~ "a_8327_8323 -> b_8327_8324",+        "t_8321_8321_8329" ~ "b_8327_8324",+        "t_8321_8321_8329" ~ "[]t_8321_8322_8321",+        "t_8321_8322_8321" ~ "i32",+        "i64" ~ "t_8321_8322_8320",+        "num_8321_8322_8326" ~ "t_8321_8322_8320",+        "t_8321_8323_8321" ~ "t_8321_8322_8320",+        "t_8321_8323_8326" ~ "t_8321_8323_8324",+        "num_8321_8323_8325" ~ "t_8321_8323_8324",+        "t_8321_8324_8323" ~ "num_8321_8324_8322",+        "num_8321_8324_8322" ~ "i64",+        "t_8321_8321_8329" ~ "[]t_8321_8324_8321",+        "t_8321_8323_8326 -> t_8321_8323_8324" ~ "a_8321_8323_8322 -> x_8321_8323_8323",+        "[]t_8321_8324_8321" ~ "[]a_8321_8323_8322",+        "t_8321_8325_8323" ~ "num_8321_8325_8322",+        "[]x_8321_8323_8323" ~ "[]update_elem_8321_8325_8325",+        "num_8321_8325_8324" ~ "update_elem_8321_8325_8325",+        "t_8321_8325_8326" ~ "[]x_8321_8323_8323",+        "t_8321_8326_8321" ~ "i32",+        "t_8321_8325_8326" ~ "[]t_8321_8326_8325",+        "{x: t_8321_8326_8321} -> (i64, num_8321_8326_8324)" ~ "a_8321_8325_8329 -> x_8321_8326_8320",+        "[]t_8321_8326_8325" ~ "[]a_8321_8325_8329",+        "[]x_8321_8326_8320" ~ "a_8321_8325_8327",+        "{xs: [](a_8321_8327_8322, b_8321_8327_8323)} -> ([]a_8321_8327_8322, []b_8321_8327_8323)" ~ "a_8321_8325_8327 -> b_8321_8325_8328",+        "(t_8321_8327_8328, t_8321_8327_8329)" ~ "b_8321_8325_8328",+        "t_8321_8328_8323 -> t_8321_8328_8323 -> t_8321_8328_8323" ~ "a_8321_8328_8322 -> a_8321_8328_8322 -> a_8321_8328_8322",+        "num_8321_8328_8324" ~ "a_8321_8328_8322",+        "t_8321_8323_8321" ~ "i64",+        "t_8321_8327_8328" ~ "[]i64",+        "t_8321_8327_8329" ~ "[]a_8321_8328_8322",+        "num_8321_8329_8327" ~ "i32",+        "t_8321_8329_8326 -> t_8321_8329_8326 -> t_8321_8329_8326" ~ "a_8321_8329_8325 -> a_8321_8329_8325 -> a_8321_8329_8325",+        "num_8321_8329_8327" ~ "a_8321_8329_8325",+        "[]a_8321_8328_8322" ~ "a_8321_8328_8320",+        "{as: []a_8321_8329_8325} -> *[]a_8321_8329_8325" ~ "a_8321_8328_8320 -> b_8321_8328_8321",+        "t_8322_8320_8326" ~ "b_8321_8328_8321",+        "t_8322_8321_8320" ~ "i64",+        "t_8322_8321_8324" ~ "i32",+        "t_8322_8321_8325" ~ "t_8322_8321_8321",+        "index_8322_8321_8326" ~ "index_elem_8322_8321_8327",+        "[]{delta_node: i32, left: #inner i32 | #leaf i32, parent: i32, right: #inner i32 | #leaf i32, sfc_code: u32}" ~ "[]index_elem_8322_8321_8327",+        "t_8322_8321_8328" ~ "index_8322_8321_8326",+        "t_8322_8321_8329" ~ "t_8322_8321_8321",+        "index_8322_8322_8320" ~ "index_elem_8322_8322_8321",+        "t_8327_8322" ~ "[]index_elem_8322_8322_8321",+        "t_8322_8322_8322" ~ "index_8322_8322_8320",+        "t_8322_8322_8323" ~ "t_8322_8321_8321",+        "index_8322_8322_8324" ~ "index_elem_8322_8322_8325",+        "t_8321_8325_8326" ~ "[]index_elem_8322_8322_8325",+        "t_8322_8322_8326" ~ "index_8322_8322_8324",+        "t_8322_8321_8328" ~ "t_8322_8323_8320",+        "kt_8322_8322_8329" ~ "t_8322_8322_8328",+        "num_8322_8323_8321" ~ "t_8322_8322_8328",+        "t_8326" ~ "t_8322_8323_8326",+        "num_8322_8323_8327" ~ "t_8322_8323_8326",+        "t_8322_8322_8328" ~ "t_8322_8322_8327",+        "t_8322_8323_8326" ~ "t_8322_8322_8327",+        "t_8322_8324_8326" ~ "t_8322_8322_8327",+        "t_8322_8322_8322" ~ "t_8322_8324_8328",+        "t_8322_8324_8329" ~ "kt_8322_8324_8327",+        "t_8322_8324_8329" ~ "t_8322_8325_8321",+        "num_8322_8325_8322" ~ "t_8322_8325_8321",+        "t_8322_8322_8326" ~ "t_8322_8325_8327",+        "t_8322_8321_8324" ~ "t_8322_8325_8327",+        "bool" ~ "t_8322_8325_8320",+        "bool" ~ "t_8322_8325_8320",+        "t_8322_8326_8326" ~ "t_8322_8325_8320",+        "t_8322_8321_8324" ~ "t_8322_8326_8327",+        "num_8322_8326_8328" ~ "t_8322_8326_8327",+        "t_8322_8322_8322" ~ "t_8322_8327_8326",+        "t_8322_8326_8326" ~ "bool",+        "t_8322_8324_8329" ~ "if_t_8322_8327_8327",+        "kt_8322_8327_8325" ~ "if_t_8322_8327_8327",+        "t_8322_8324_8326" ~ "t_8322_8327_8324",+        "if_t_8322_8327_8327" ~ "t_8322_8327_8324",+        "bool" ~ "bool",+        "num_8322_8327_8323" ~ "if_t_8322_8328_8322",+        "t_8322_8327_8324" ~ "if_t_8322_8328_8322",+        "t_8322_8328_8323" ~ "if_t_8322_8328_8322",+        "t_8322_8321_8328" ~ "t_8322_8328_8325",+        "t_8322_8321_8328" ~ "t_8322_8328_8327",+        "t_8322_8326_8326" ~ "bool",+        "kt_8322_8328_8324" ~ "if_t_8322_8328_8328",+        "kt_8322_8328_8326" ~ "if_t_8322_8328_8328",+        "t_8322_8328_8329" ~ "if_t_8322_8328_8328",+        "t_8322_8321_8324" ~ "t_8322_8329_8321",+        "num_8322_8329_8322" ~ "t_8322_8329_8321",+        "t_8322_8329_8328" ~ "t_8322_8328_8329",+        "t_8323_8320_8320" ~ "t_8322_8328_8329",+        "match_t_8323_8320_8321" ~ "bool",+        "bool" ~ "t_8322_8329_8320",+        "match_t_8323_8320_8321" ~ "t_8322_8329_8320",+        "t_8323_8320_8326" ~ "t_8322_8329_8320",+        "t_8323_8320_8329" ~ "t_8323_8320_8327",+        "t_8323_8321_8321" ~ "t_8323_8320_8327",+        "t_8323_8320_8328" ~ "t_8323_8321_8320",+        "match_t_8323_8321_8322" ~ "t_8323_8320_8328",+        "i32" ~ "match_t_8323_8321_8322",+        "t_8322_8328_8329" ~ "t_8323_8320_8327",+        "t_8323_8321_8325" ~ "i32",+        "index_8323_8321_8326" ~ "index_elem_8323_8321_8327",+        "[]{mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}" ~ "[]index_elem_8323_8321_8327",+        "t_8323_8321_8328" ~ "index_8323_8321_8326",+        "t_8323_8321_8328" ~ "t_8323_8322_8320",+        "t_8323_8321_8328" ~ "t_8323_8322_8322",+        "kt_8323_8321_8329" ~ "f32",+        "kt_8323_8322_8321" ~ "{x: f32, y: f32, z: f32}",+        "t_8323_8321_8328" ~ "ft_8323_8322_8327",+        "(num_8323_8322_8328, num_8323_8322_8329, num_8323_8323_8320)" ~ "(f32, f32, f32)",+        "(num_8323_8323_8324, num_8323_8323_8325, num_8323_8323_8326)" ~ "(f32, f32, f32)",+        "t_8323_8320_8326" ~ "bool",+        "t_8323_8321_8328" ~ "if_t_8323_8323_8329",+        "{mass: float_8323_8323_8323, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}" ~ "if_t_8323_8323_8329",+        "t_8323_8324_8320" ~ "if_t_8323_8323_8329",+        "t_8322_8321_8324" ~ "t_8323_8324_8321",+        "num_8323_8324_8322" ~ "t_8323_8324_8321",+        "t_8323_8324_8328" ~ "num_8323_8324_8327",+        "t_8323_8324_8329" ~ "t_8322_8321_8321",+        "t_8323_8325_8321" ~ "num_8323_8325_8320",+        "t_8323_8325_8322" ~ "num_8323_8325_8320",+        "(t_8323_8325_8323, t_8323_8325_8324)" ~ "(t_8323_8325_8322, t_8323_8324_8329)",+        "t_8323_8325_8327" ~ "num_8323_8325_8326",+        "t_8323_8325_8323" ~ "t_8323_8325_8325",+        "num_8323_8325_8326" ~ "t_8323_8325_8325",+        "t_8323_8326_8322" ~ "t_8323_8325_8324",+        "index_8323_8326_8323" ~ "index_elem_8323_8326_8324",+        "[]{delta_node: i32, left: #inner i32 | #leaf i32, parent: i32, right: #inner i32 | #leaf i32, sfc_code: u32}" ~ "[]index_elem_8323_8326_8324",+        "index_8323_8326_8323" ~ "t_8323_8326_8326",+        "t_8323_8326_8327" ~ "kt_8323_8326_8325",+        "t_8323_8327_8320" ~ "num_8323_8326_8329",+        "t_8323_8326_8327" ~ "t_8323_8326_8328",+        "num_8323_8326_8329" ~ "t_8323_8326_8328",+        "t_8323_8327_8326" ~ "t_8323_8326_8327",+        "index_8323_8327_8327" ~ "index_elem_8323_8327_8328",+        "[]{delta_node: i32, left: #inner i32 | #leaf i32, parent: i32, right: #inner i32 | #leaf i32, sfc_code: u32}" ~ "[]index_elem_8323_8327_8328",+        "index_8323_8327_8327" ~ "t_8323_8328_8320",+        "kt_8323_8327_8329" ~ "t_8323_8320_8327",+        "t_8323_8328_8323" ~ "i32",+        "t_8323_8328_8324" ~ "t_8323_8326_8327",+        "index_8323_8328_8325" ~ "index_elem_8323_8328_8326",+        "t_8327_8322" ~ "[]index_elem_8323_8328_8326",+        "t_8323_8328_8327" ~ "index_8323_8328_8325",+        "t_8323_8328_8323" ~ "t_8323_8328_8328",+        "t_8323_8325_8324" ~ "t_8323_8328_8328",+        "t_8323_8329_8323" ~ "bool",+        "t_8323_8328_8327" ~ "t_8323_8329_8325",+        "t_8323_8328_8327" ~ "t_8323_8329_8327",+        "t_8323_8329_8323" ~ "bool",+        "kt_8323_8329_8324" ~ "if_t_8323_8329_8328",+        "kt_8323_8329_8326" ~ "if_t_8323_8329_8328",+        "t_8323_8329_8329" ~ "if_t_8323_8329_8328",+        "t_8323_8329_8329" ~ "t_8324_8320_8320",+        "num_8324_8320_8321" ~ "t_8324_8320_8320",+        "t_8324_8320_8326" ~ "t_8323_8326_8327",+        "index_8324_8320_8327" ~ "index_elem_8324_8320_8328",+        "t_8321_8325_8326" ~ "[]index_elem_8324_8320_8328",+        "t_8324_8320_8329" ~ "index_8324_8320_8327",+        "t_8323_8328_8327" ~ "t_8324_8321_8325",+        "kt_8324_8321_8324" ~ "t_8324_8321_8323",+        "num_8324_8321_8326" ~ "t_8324_8321_8323",+        "bool" ~ "bool",+        "num_8324_8321_8322" ~ "t_8324_8321_8321",+        "i32" ~ "t_8324_8321_8321",+        "t_8324_8320_8329" ~ "t_8324_8321_8320",+        "t_8324_8321_8321" ~ "t_8324_8321_8320",+        "t_8323_8329_8323" ~ "bool",+        "t_8324_8320_8329" ~ "if_t_8324_8323_8321",+        "t_8324_8321_8320" ~ "if_t_8324_8323_8321",+        "bool" ~ "bool",+        "(if_t_8324_8323_8321, t_8323_8326_8327)" ~ "if_t_8324_8323_8322",+        "(t_8323_8325_8323, t_8323_8326_8327)" ~ "if_t_8324_8323_8322",+        "bool" ~ "bool",+        "(num_8323_8327_8325, t_8323_8326_8327)" ~ "if_t_8324_8323_8323",+        "if_t_8324_8323_8322" ~ "if_t_8324_8323_8323",+        "(t_8323_8325_8322, t_8323_8324_8329)" ~ "if_t_8324_8323_8323",+        "(t_8324_8323_8324, t_8324_8323_8325)" ~ "(t_8323_8325_8323, t_8323_8325_8324)",+        "bool" ~ "bool",+        "num_8323_8324_8327" ~ "if_t_8324_8323_8326",+        "t_8324_8323_8324" ~ "if_t_8324_8323_8326",+        "t_8324_8323_8327" ~ "if_t_8324_8323_8326",+        "t_8324_8323_8329" ~ "num_8324_8323_8328",+        "i32" ~ "num_8324_8323_8328",+        "num_8324_8324_8321" ~ "i64",+        "i32" ~ "t_8324_8324_8320",+        "t_8324_8324_8326" ~ "[]t_8324_8324_8320",+        "t_8322_8321_8324" ~ "t_8324_8324_8327",+        "num_8324_8324_8328" ~ "t_8324_8324_8327",+        "t_8322_8328_8329" ~ "t_8323_8320_8327",+        "bool" ~ "bool",+        "num_8324_8325_8323" ~ "if_t_8324_8325_8326",+        "i32" ~ "if_t_8324_8325_8326",+        "t_8324_8325_8327" ~ "if_t_8324_8325_8326",+        "i32" ~ "num_8324_8325_8328",+        "num_8324_8326_8320" ~ "et_8324_8325_8329",+        "num_8324_8326_8321" ~ "et_8324_8325_8329",+        "num_8324_8326_8322" ~ "et_8324_8325_8329",+        "t_8324_8326_8323" ~ "[]et_8324_8325_8329",+        "i32" ~ "num_8324_8326_8324",+        "(t_8324_8326_8325, t_8324_8326_8326, t_8324_8326_8327, t_8324_8326_8328, t_8324_8326_8329)" ~ "(t_8324_8324_8326, t_8324_8326_8323, i32, t_8324_8325_8327, i32)",+        "t_8324_8327_8322" ~ "num_8324_8327_8321",+        "index_8324_8327_8323" ~ "index_elem_8324_8327_8324",+        "t_8324_8326_8326" ~ "[]index_elem_8324_8327_8324",+        "index_8324_8327_8323" ~ "t_8324_8327_8320",+        "num_8324_8327_8325" ~ "t_8324_8327_8320",+        "t_8324_8328_8321" ~ "t_8324_8326_8329",+        "index_8324_8328_8322" ~ "index_elem_8324_8328_8323",+        "t_8324_8326_8326" ~ "[]index_elem_8324_8328_8323",+        "index_8324_8328_8322" ~ "t_8324_8328_8320",+        "num_8324_8328_8324" ~ "t_8324_8328_8320",+        "t_8324_8329_8320" ~ "t_8324_8326_8328",+        "index_8324_8329_8321" ~ "index_elem_8324_8329_8322",+        "t_8327_8322" ~ "[]index_elem_8324_8329_8322",+        "index_8324_8329_8321" ~ "t_8324_8329_8324",+        "kt_8324_8329_8323" ~ "t_8324_8328_8329",+        "num_8324_8329_8325" ~ "t_8324_8328_8329",+        "t_8325_8320_8320" ~ "t_8324_8326_8328",+        "index_8325_8320_8321" ~ "index_elem_8325_8320_8322",+        "t_8321_8325_8326" ~ "[]index_elem_8325_8320_8322",+        "t_8325_8320_8323" ~ "index_8325_8320_8321",+        "t_8325_8320_8324" ~ "t_8324_8326_8327",+        "t_8324_8326_8325" ~ "[]update_elem_8325_8320_8325",+        "t_8325_8320_8323" ~ "update_elem_8325_8320_8325",+        "t_8324_8326_8327" ~ "t_8325_8320_8326",+        "num_8325_8320_8327" ~ "t_8325_8320_8326",+        "t_8325_8321_8322" ~ "t_8325_8320_8326",+        "t_8325_8321_8323" ~ "t_8324_8326_8329",+        "t_8324_8326_8326" ~ "[]update_elem_8325_8321_8325",+        "num_8325_8321_8324" ~ "update_elem_8325_8321_8325",+        "t_8325_8321_8326" ~ "t_8324_8326_8328",+        "index_8325_8321_8327" ~ "index_elem_8325_8321_8328",+        "[]{delta_node: i32, left: #inner i32 | #leaf i32, parent: i32, right: #inner i32 | #leaf i32, sfc_code: u32}" ~ "[]index_elem_8325_8321_8328",+        "index_8325_8321_8327" ~ "t_8325_8322_8320",+        "kt_8325_8321_8329" ~ "t_8323_8320_8327",+        "t_8325_8322_8323" ~ "i32",+        "t_8324_8326_8329" ~ "t_8325_8322_8324",+        "num_8325_8322_8325" ~ "t_8325_8322_8324",+        "t_8325_8323_8320" ~ "t_8325_8322_8324",+        "t_8325_8323_8321" ~ "t_8325_8323_8320",+        "t_8324_8326_8326" ~ "[]update_elem_8325_8323_8323",+        "num_8325_8323_8322" ~ "update_elem_8325_8323_8323",+        "bool" ~ "bool",+        "(t_8324_8326_8325, t_8324_8326_8326, t_8325_8321_8322, t_8324_8326_8328, t_8324_8326_8329)" ~ "if_t_8325_8323_8324",+        "(t_8324_8326_8325, t_8324_8326_8326, t_8324_8326_8327, t_8325_8322_8323, t_8325_8323_8320)" ~ "if_t_8325_8323_8324",+        "t_8325_8323_8326" ~ "t_8324_8326_8329",+        "index_8325_8323_8327" ~ "index_elem_8325_8323_8328",+        "t_8324_8326_8326" ~ "[]index_elem_8325_8323_8328",+        "index_8325_8323_8327" ~ "t_8325_8323_8325",+        "num_8325_8323_8329" ~ "t_8325_8323_8325",+        "t_8325_8324_8325" ~ "t_8324_8326_8328",+        "index_8325_8324_8326" ~ "index_elem_8325_8324_8327",+        "t_8327_8322" ~ "[]index_elem_8325_8324_8327",+        "index_8325_8324_8326" ~ "t_8325_8324_8329",+        "kt_8325_8324_8328" ~ "t_8325_8324_8324",+        "num_8325_8325_8320" ~ "t_8325_8324_8324",+        "t_8325_8325_8325" ~ "t_8324_8326_8328",+        "index_8325_8325_8326" ~ "index_elem_8325_8325_8327",+        "t_8321_8325_8326" ~ "[]index_elem_8325_8325_8327",+        "t_8325_8325_8328" ~ "index_8325_8325_8326",+        "t_8325_8325_8329" ~ "t_8324_8326_8327",+        "t_8325_8326_8323" ~ "t_8324_8326_8328",+        "index_8325_8326_8324" ~ "index_elem_8325_8326_8325",+        "t_8327_8322" ~ "[]index_elem_8325_8326_8325",+        "index_8325_8326_8324" ~ "t_8325_8326_8327",+        "kt_8325_8326_8326" ~ "t_8325_8326_8322",+        "num_8325_8326_8328" ~ "t_8325_8326_8322",+        "bool" ~ "bool",+        "i32" ~ "t_8325_8326_8321",+        "num_8325_8327_8325" ~ "t_8325_8326_8321",+        "t_8325_8325_8328" ~ "t_8325_8326_8320",+        "t_8325_8326_8321" ~ "t_8325_8326_8320",+        "t_8324_8326_8325" ~ "[]update_elem_8325_8328_8324",+        "t_8325_8326_8320" ~ "update_elem_8325_8328_8324",+        "t_8324_8326_8327" ~ "t_8325_8328_8325",+        "num_8325_8328_8326" ~ "t_8325_8328_8325",+        "t_8325_8329_8321" ~ "t_8325_8328_8325",+        "t_8325_8329_8322" ~ "t_8324_8326_8329",+        "t_8324_8326_8326" ~ "[]update_elem_8325_8329_8324",+        "num_8325_8329_8323" ~ "update_elem_8325_8329_8324",+        "t_8325_8329_8325" ~ "t_8324_8326_8328",+        "index_8325_8329_8326" ~ "index_elem_8325_8329_8327",+        "[]{delta_node: i32, left: #inner i32 | #leaf i32, parent: i32, right: #inner i32 | #leaf i32, sfc_code: u32}" ~ "[]index_elem_8325_8329_8327",+        "index_8325_8329_8326" ~ "t_8325_8329_8329",+        "kt_8325_8329_8328" ~ "t_8323_8320_8327",+        "t_8326_8320_8322" ~ "i32",+        "t_8324_8326_8329" ~ "t_8326_8320_8323",+        "num_8326_8320_8324" ~ "t_8326_8320_8323",+        "t_8326_8320_8329" ~ "t_8326_8320_8323",+        "t_8326_8321_8320" ~ "t_8326_8320_8329",+        "t_8324_8326_8326" ~ "[]update_elem_8326_8321_8322",+        "num_8326_8321_8321" ~ "update_elem_8326_8321_8322",+        "bool" ~ "bool",+        "(t_8324_8326_8325, t_8324_8326_8326, t_8325_8329_8321, t_8324_8326_8328, t_8324_8326_8329)" ~ "if_t_8326_8321_8323",+        "(t_8324_8326_8325, t_8324_8326_8326, t_8324_8326_8327, t_8326_8320_8322, t_8326_8320_8329)" ~ "if_t_8326_8321_8323",+        "t_8326_8321_8324" ~ "t_8324_8326_8328",+        "index_8326_8321_8325" ~ "index_elem_8326_8321_8326",+        "[]{delta_node: i32, left: #inner i32 | #leaf i32, parent: i32, right: #inner i32 | #leaf i32, sfc_code: u32}" ~ "[]index_elem_8326_8321_8326",+        "index_8326_8321_8325" ~ "t_8326_8321_8328",+        "t_8326_8321_8329" ~ "kt_8326_8321_8327",+        "t_8324_8326_8329" ~ "t_8326_8322_8320",+        "num_8326_8322_8321" ~ "t_8326_8322_8320",+        "t_8326_8322_8326" ~ "t_8326_8322_8320",+        "t_8326_8322_8327" ~ "t_8326_8322_8326",+        "t_8326_8322_8329" ~ "t_8326_8322_8326",+        "index_8326_8323_8320" ~ "index_elem_8326_8323_8321",+        "t_8324_8326_8326" ~ "[]index_elem_8326_8323_8321",+        "num_8326_8323_8322" ~ "i32",+        "index_8326_8323_8320" ~ "t_8326_8322_8328",+        "num_8326_8323_8322" ~ "t_8326_8322_8328",+        "t_8324_8326_8326" ~ "[]update_elem_8326_8323_8327",+        "t_8326_8322_8328" ~ "update_elem_8326_8323_8327",+        "bool" ~ "bool",+        "if_t_8326_8321_8323" ~ "if_t_8326_8323_8328",+        "(t_8324_8326_8325, t_8324_8326_8326, t_8324_8326_8327, t_8326_8321_8329, t_8326_8322_8326)" ~ "if_t_8326_8323_8328",+        "bool" ~ "bool",+        "if_t_8325_8323_8324" ~ "if_t_8326_8323_8329",+        "if_t_8326_8323_8328" ~ "if_t_8326_8323_8329",+        "(t_8324_8324_8326, t_8324_8326_8323, i32, t_8324_8325_8327, i32)" ~ "if_t_8326_8323_8329",+        "(t_8326_8324_8320, t_8326_8324_8321, t_8326_8324_8322, t_8326_8324_8323, t_8326_8324_8324)" ~ "(t_8324_8326_8325, t_8324_8326_8326, t_8324_8326_8327, t_8324_8326_8328, t_8324_8326_8329)",+        "t_8323_8320_8326" ~ "bool",+        "t_8324_8324_8326" ~ "if_t_8326_8324_8325",+        "t_8326_8324_8320" ~ "if_t_8326_8324_8325",+        "t_8326_8324_8326" ~ "if_t_8326_8324_8325",+        "t_8321_8323_8321" ~ "i64",+        "{i: t_8322_8321_8320} -> {rp: t_8322_8321_8321} -> {body: t_8323_8324_8320, children: t_8326_8324_8326, is_leaf: t_8323_8320_8326, parent: t_8324_8323_8327, tree_level: t_8322_8328_8323}" ~ "a_8322_8320_8327 -> b_8322_8320_8328 -> x_8322_8320_8329",+        "[]i64" ~ "[]a_8322_8320_8327",+        "t_8322_8320_8326" ~ "[]b_8322_8320_8328",+        "t_8326_8325_8325" ~ "[]x_8322_8320_8329",+        "[]{body: {mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}, children: []i32, is_leaf: bool, parent: i32, tree_level: i32}" ~ "t_8326_8325_8325"+      ],+      M.empty,+      M.fromList [("num_8320", (3, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321", (3, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8322", (3, TyVarFree NoLoc Unlifted)), ("index_elem_8323", (3, TyVarFree NoLoc Unlifted)), ("kt_8324", (3, TyVarFree NoLoc Lifted)), ("t_8325", (3, TyVarRecord NoLoc (M.fromList [("delta_node", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8324" 15206}) []))]))), ("t_8326", (4, TyVarFree NoLoc Lifted)), ("t_8327", (6, TyVarFree NoLoc Lifted)), ("t_8328", (7, TyVarFree NoLoc Lifted)), ("t_8329", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8320", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8321", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8326", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8327", (8, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8326", (6, TyVarFree NoLoc Lifted)), ("t_8322_8327", (7, TyVarFree NoLoc Lifted)), ("t_8322_8328", (7, TyVarSum NoLoc (M.fromList [("leaf", [Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "t_8322_8327" 15232}) [])])]))), ("t_8322_8329", (7, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8320", (7, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8325", (7, TyVarFree NoLoc Lifted)), ("t_8323_8326", (7, TyVarSum NoLoc (M.fromList [("inner", [Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "t_8323_8325" 15241}) [])])]))), ("t_8323_8327", (7, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8323_8328", (7, TyVarFree NoLoc Unlifted)), ("index_elem_8323_8329", (7, TyVarFree NoLoc Unlifted)), ("kt_8324_8320", (7, TyVarFree NoLoc Lifted)), ("t_8324_8321", (7, TyVarRecord NoLoc (M.fromList [("delta_node", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8324_8320" 15247}) []))]))), ("match_t_8324_8322", (7, TyVarFree NoLoc SizeLifted)), ("a_8324_8323", (5, TyVarFree NoLoc Unlifted)), ("x_8324_8324", (5, TyVarFree NoLoc Unlifted)), ("t_8324_8325", (6, TyVarFree NoLoc Lifted)), ("kt_8324_8326", (7, TyVarFree NoLoc Lifted)), ("t_8324_8327", (7, TyVarRecord NoLoc (M.fromList [("left", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8324_8326" 15255}) []))]))), ("kt_8325_8320", (7, TyVarFree NoLoc Lifted)), ("t_8325_8321", (7, TyVarRecord NoLoc (M.fromList [("delta_node", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8325_8320" 15260}) []))]))), ("t_8325_8326", (8, TyVarFree NoLoc Lifted)), ("kt_8325_8327", (9, TyVarFree NoLoc Lifted)), ("t_8325_8328", (9, TyVarRecord NoLoc (M.fromList [("right", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8325_8327" 15268}) []))]))), ("kt_8326_8321", (9, TyVarFree NoLoc Lifted)), ("t_8326_8322", (9, TyVarRecord NoLoc (M.fromList [("delta_node", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8326_8321" 15273}) []))]))), ("t_8326_8327", (10, TyVarFree NoLoc Lifted)), ("t_8327_8322", (6, TyVarFree NoLoc Lifted)), ("a_8327_8323", (7, TyVarFree NoLoc Lifted)), ("b_8327_8324", (7, TyVarFree NoLoc Lifted)), ("a_8327_8325", (7, TyVarFree NoLoc Lifted)), ("b_8327_8326", (7, TyVarFree NoLoc Lifted)), ("a_8327_8327", (7, TyVarFree NoLoc Unlifted)), ("x_8327_8328", (7, TyVarFree NoLoc Unlifted)), ("t_8327_8329", (8, TyVarFree NoLoc Lifted)), ("t_8328_8320", (8, TyVarFree NoLoc Lifted)), ("t_8328_8321", (9, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8328_8322", (9, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8328_8323", (9, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8329_8320", (9, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8329_8321", (9, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("a_8321_8320_8328", (7, TyVarFree NoLoc Unlifted)), ("t_8321_8320_8329", (7, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8321_8320", (7, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8321_8329", (8, TyVarFree NoLoc Lifted)), ("t_8321_8322_8320", (9, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8322_8321", (9, TyVarFree NoLoc Unlifted)), ("num_8321_8322_8326", (9, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8323_8321", (10, TyVarFree NoLoc Lifted)), ("a_8321_8323_8322", (11, TyVarFree NoLoc Unlifted)), ("x_8321_8323_8323", (11, TyVarFree NoLoc Unlifted)), ("t_8321_8323_8324", (11, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8323_8325", (11, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8323_8326", (11, TyVarFree NoLoc Lifted)), ("t_8321_8324_8321", (11, TyVarFree NoLoc Unlifted)), ("num_8321_8324_8322", (11, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8324_8323", (11, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8325_8322", (11, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8325_8323", (11, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("num_8321_8325_8324", (11, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("update_elem_8321_8325_8325", (11, TyVarFree NoLoc Unlifted)), ("t_8321_8325_8326", (12, TyVarFree NoLoc Lifted)), ("a_8321_8325_8327", (13, TyVarFree NoLoc Lifted)), ("b_8321_8325_8328", (13, TyVarFree NoLoc Lifted)), ("a_8321_8325_8329", (13, TyVarFree NoLoc Unlifted)), ("x_8321_8326_8320", (13, TyVarFree NoLoc Unlifted)), ("t_8321_8326_8321", (14, TyVarFree NoLoc Lifted)), ("num_8321_8326_8324", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8326_8325", (13, TyVarFree NoLoc Unlifted)), ("a_8321_8327_8322", (13, TyVarFree NoLoc Unlifted)), ("b_8321_8327_8323", (13, TyVarFree NoLoc Unlifted)), ("t_8321_8327_8328", (14, TyVarFree NoLoc Lifted)), ("t_8321_8327_8329", (14, TyVarFree NoLoc Lifted)), ("a_8321_8328_8320", (15, TyVarFree NoLoc Lifted)), ("b_8321_8328_8321", (15, TyVarFree NoLoc Lifted)), ("a_8321_8328_8322", (15, TyVarFree NoLoc Unlifted)), ("t_8321_8328_8323", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8328_8324", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("a_8321_8329_8325", (15, TyVarFree NoLoc Unlifted)), ("t_8321_8329_8326", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8329_8327", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8320_8326", (16, TyVarFree NoLoc Lifted)), ("a_8322_8320_8327", (17, TyVarFree NoLoc Unlifted)), ("b_8322_8320_8328", (17, TyVarFree NoLoc Unlifted)), ("x_8322_8320_8329", (17, TyVarFree NoLoc Unlifted)), ("t_8322_8321_8320", (18, TyVarFree NoLoc Lifted)), ("t_8322_8321_8321", (19, TyVarFree NoLoc Lifted)), ("t_8322_8321_8324", (21, TyVarFree NoLoc Lifted)), ("t_8322_8321_8325", (22, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8322_8321_8326", (22, TyVarFree NoLoc Unlifted)), ("index_elem_8322_8321_8327", (22, TyVarFree NoLoc Unlifted)), ("t_8322_8321_8328", (23, TyVarFree NoLoc Lifted)), ("t_8322_8321_8329", (24, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8322_8322_8320", (24, TyVarFree NoLoc Unlifted)), ("index_elem_8322_8322_8321", (24, TyVarFree NoLoc Unlifted)), ("t_8322_8322_8322", (25, TyVarFree NoLoc Lifted)), ("t_8322_8322_8323", (26, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8322_8322_8324", (26, TyVarFree NoLoc Unlifted)), ("index_elem_8322_8322_8325", (26, TyVarFree NoLoc Unlifted)), ("t_8322_8322_8326", (27, TyVarFree NoLoc Lifted)), ("t_8322_8322_8327", (28, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8322_8328", (28, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("kt_8322_8322_8329", (28, TyVarFree NoLoc Lifted)), ("t_8322_8323_8320", (28, TyVarRecord NoLoc (M.fromList [("delta_node", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8322_8322_8329" 15494}) []))]))), ("num_8322_8323_8321", (28, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8323_8326", (28, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8323_8327", (28, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8324_8326", (29, TyVarFree NoLoc Lifted)), ("kt_8322_8324_8327", (30, TyVarFree NoLoc Lifted)), ("t_8322_8324_8328", (30, TyVarRecord NoLoc (M.fromList [("left", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8322_8324_8327" 15515}) []))]))), ("t_8322_8324_8329", (31, TyVarFree NoLoc Lifted)), ("t_8322_8325_8320", (32, TyVarPrim NoLoc [Bool])), ("t_8322_8325_8321", (32, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8325_8322", (32, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8325_8327", (32, TyVarFree NoLoc Unlifted)), ("t_8322_8326_8326", (33, TyVarFree NoLoc Lifted)), ("t_8322_8326_8327", (34, TyVarFree NoLoc Unlifted)), ("num_8322_8326_8328", (34, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8327_8323", (34, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8327_8324", (34, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("kt_8322_8327_8325", (34, TyVarFree NoLoc Lifted)), ("t_8322_8327_8326", (34, TyVarRecord NoLoc (M.fromList [("right", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8322_8327_8325" 15547}) []))]))), ("if_t_8322_8327_8327", (34, TyVarFree NoLoc SizeLifted)), ("if_t_8322_8328_8322", (34, TyVarFree NoLoc SizeLifted)), ("t_8322_8328_8323", (35, TyVarFree NoLoc Lifted)), ("kt_8322_8328_8324", (36, TyVarFree NoLoc Lifted)), ("t_8322_8328_8325", (36, TyVarRecord NoLoc (M.fromList [("left", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8322_8328_8324" 15557}) []))]))), ("kt_8322_8328_8326", (36, TyVarFree NoLoc Lifted)), ("t_8322_8328_8327", (36, TyVarRecord NoLoc (M.fromList [("right", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8322_8328_8326" 15559}) []))]))), ("if_t_8322_8328_8328", (36, TyVarFree NoLoc SizeLifted)), ("t_8322_8328_8329", (37, TyVarFree NoLoc Lifted)), ("t_8322_8329_8320", (38, TyVarPrim NoLoc [Bool])), ("t_8322_8329_8321", (38, TyVarFree NoLoc Unlifted)), ("num_8322_8329_8322", (38, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8329_8327", (38, TyVarFree NoLoc Lifted)), ("t_8322_8329_8328", (38, TyVarSum NoLoc (M.fromList [("leaf", [Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "t_8322_8329_8327" 15571}) [])])]))), ("t_8322_8329_8329", (38, TyVarFree NoLoc Lifted)), ("t_8323_8320_8320", (38, TyVarSum NoLoc (M.fromList [("inner", [Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "t_8322_8329_8329" 15573}) [])])]))), ("match_t_8323_8320_8321", (38, TyVarFree NoLoc SizeLifted)), ("t_8323_8320_8326", (39, TyVarFree NoLoc Lifted)), ("t_8323_8320_8327", (41, TyVarFree NoLoc Lifted)), ("t_8323_8320_8328", (42, TyVarFree NoLoc Lifted)), ("t_8323_8320_8329", (42, TyVarSum NoLoc (M.fromList [("inner", [Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "t_8323_8320_8328" 15583}) [])])]))), ("t_8323_8321_8320", (42, TyVarFree NoLoc Lifted)), ("t_8323_8321_8321", (42, TyVarSum NoLoc (M.fromList [("leaf", [Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "t_8323_8321_8320" 15585}) [])])]))), ("match_t_8323_8321_8322", (42, TyVarFree NoLoc SizeLifted)), ("t_8323_8321_8325", (40, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8323_8321_8326", (40, TyVarFree NoLoc Unlifted)), ("index_elem_8323_8321_8327", (40, TyVarFree NoLoc Unlifted)), ("t_8323_8321_8328", (41, TyVarFree NoLoc Lifted)), ("kt_8323_8321_8329", (42, TyVarFree NoLoc Lifted)), ("t_8323_8322_8320", (42, TyVarRecord NoLoc (M.fromList [("mass", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8323_8321_8329" 15596}) []))]))), ("kt_8323_8322_8321", (42, TyVarFree NoLoc Lifted)), ("t_8323_8322_8322", (42, TyVarRecord NoLoc (M.fromList [("position", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8323_8322_8321" 15598}) []))]))), ("ft_8323_8322_8327", (42, TyVarRecord NoLoc (M.fromList [("position", Scalar (Record (M.fromList [("x", Scalar (Prim (FloatType Float32))), ("y", Scalar (Prim (FloatType Float32))), ("z", Scalar (Prim (FloatType Float32)))])))]))), ("num_8323_8322_8328", (40, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8322_8329", (40, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8323_8320", (40, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("float_8323_8323_8323", (40, TyVarPrim NoLoc [FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8323_8324", (40, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8323_8325", (40, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8323_8326", (40, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("if_t_8323_8323_8329", (40, TyVarFree NoLoc SizeLifted)), ("t_8323_8324_8320", (41, TyVarFree NoLoc Lifted)), ("t_8323_8324_8321", (42, TyVarFree NoLoc Unlifted)), ("num_8323_8324_8322", (42, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8324_8327", (42, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8324_8328", (42, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8324_8329", (43, TyVarFree NoLoc Lifted)), ("num_8323_8325_8320", (44, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8325_8321", (44, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8325_8322", (45, TyVarFree NoLoc Lifted)), ("t_8323_8325_8323", (46, TyVarFree NoLoc Lifted)), ("t_8323_8325_8324", (46, TyVarFree NoLoc Lifted)), ("t_8323_8325_8325", (46, TyVarFree NoLoc Unlifted)), ("num_8323_8325_8326", (46, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8325_8327", (46, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8326_8322", (46, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8323_8326_8323", (46, TyVarFree NoLoc Unlifted)), ("index_elem_8323_8326_8324", (46, TyVarFree NoLoc Unlifted)), ("kt_8323_8326_8325", (46, TyVarFree NoLoc Lifted)), ("t_8323_8326_8326", (46, TyVarRecord NoLoc (M.fromList [("parent", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8323_8326_8325" 15648}) []))]))), ("t_8323_8326_8327", (47, TyVarFree NoLoc Lifted)), ("t_8323_8326_8328", (48, TyVarFree NoLoc Unlifted)), ("num_8323_8326_8329", (48, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8327_8320", (48, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8327_8325", (48, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8327_8326", (48, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8323_8327_8327", (48, TyVarFree NoLoc Unlifted)), ("index_elem_8323_8327_8328", (48, TyVarFree NoLoc Unlifted)), ("kt_8323_8327_8329", (48, TyVarFree NoLoc Lifted)), ("t_8323_8328_8320", (48, TyVarRecord NoLoc (M.fromList [("left", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8323_8327_8329" 15664}) []))]))), ("t_8323_8328_8323", (49, TyVarFree NoLoc Lifted)), ("t_8323_8328_8324", (50, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8323_8328_8325", (50, TyVarFree NoLoc Unlifted)), ("index_elem_8323_8328_8326", (50, TyVarFree NoLoc Unlifted)), ("t_8323_8328_8327", (51, TyVarFree NoLoc Lifted)), ("t_8323_8328_8328", (52, TyVarFree NoLoc Unlifted)), ("t_8323_8329_8323", (53, TyVarFree NoLoc Lifted)), ("kt_8323_8329_8324", (54, TyVarFree NoLoc Lifted)), ("t_8323_8329_8325", (54, TyVarRecord NoLoc (M.fromList [("left", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8323_8329_8324" 15681}) []))]))), ("kt_8323_8329_8326", (54, TyVarFree NoLoc Lifted)), ("t_8323_8329_8327", (54, TyVarRecord NoLoc (M.fromList [("right", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8323_8329_8326" 15683}) []))]))), ("if_t_8323_8329_8328", (54, TyVarFree NoLoc SizeLifted)), ("t_8323_8329_8329", (55, TyVarFree NoLoc Lifted)), ("t_8324_8320_8320", (56, TyVarFree NoLoc Unlifted)), ("num_8324_8320_8321", (56, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8324_8320_8326", (56, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8324_8320_8327", (56, TyVarFree NoLoc Unlifted)), ("index_elem_8324_8320_8328", (56, TyVarFree NoLoc Unlifted)), ("t_8324_8320_8329", (57, TyVarFree NoLoc Lifted)), ("t_8324_8321_8320", (58, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8324_8321_8321", (58, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8324_8321_8322", (58, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8324_8321_8323", (58, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("kt_8324_8321_8324", (58, TyVarFree NoLoc Lifted)), ("t_8324_8321_8325", (58, TyVarRecord NoLoc (M.fromList [("left", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8324_8321_8324" 15702}) []))]))), ("num_8324_8321_8326", (58, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("if_t_8324_8323_8321", (58, TyVarFree NoLoc SizeLifted)), ("if_t_8324_8323_8322", (56, TyVarFree NoLoc SizeLifted)), ("if_t_8324_8323_8323", (48, TyVarFree NoLoc SizeLifted)), ("t_8324_8323_8324", (47, TyVarFree NoLoc Lifted)), ("t_8324_8323_8325", (47, TyVarFree NoLoc Lifted)), ("if_t_8324_8323_8326", (42, TyVarFree NoLoc SizeLifted)), ("t_8324_8323_8327", (43, TyVarFree NoLoc Lifted)), ("num_8324_8323_8328", (44, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8324_8323_8329", (44, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8324_8324_8320", (46, TyVarFree NoLoc Unlifted)), ("num_8324_8324_8321", (46, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8324_8324_8326", (47, TyVarFree NoLoc Lifted)), ("t_8324_8324_8327", (48, TyVarFree NoLoc Unlifted)), ("num_8324_8324_8328", (48, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8324_8325_8323", (48, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("if_t_8324_8325_8326", (48, TyVarFree NoLoc SizeLifted)), ("t_8324_8325_8327", (49, TyVarFree NoLoc Lifted)), ("num_8324_8325_8328", (50, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("et_8324_8325_8329", (52, TyVarFree NoLoc Unlifted)), ("num_8324_8326_8320", (52, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8324_8326_8321", (52, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8324_8326_8322", (52, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8324_8326_8323", (53, TyVarFree NoLoc Lifted)), ("num_8324_8326_8324", (54, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8324_8326_8325", (56, TyVarFree NoLoc Lifted)), ("t_8324_8326_8326", (56, TyVarFree NoLoc Lifted)), ("t_8324_8326_8327", (56, TyVarFree NoLoc Lifted)), ("t_8324_8326_8328", (56, TyVarFree NoLoc Lifted)), ("t_8324_8326_8329", (56, TyVarFree NoLoc Lifted)), ("t_8324_8327_8320", (56, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8324_8327_8321", (56, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8324_8327_8322", (56, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8324_8327_8323", (56, TyVarFree NoLoc Unlifted)), ("index_elem_8324_8327_8324", (56, TyVarFree NoLoc Unlifted)), ("num_8324_8327_8325", (56, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8324_8328_8320", (56, TyVarFree NoLoc Unlifted)), ("t_8324_8328_8321", (56, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8324_8328_8322", (56, TyVarFree NoLoc Unlifted)), ("index_elem_8324_8328_8323", (56, TyVarFree NoLoc Unlifted)), ("num_8324_8328_8324", (56, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8324_8328_8329", (56, TyVarFree NoLoc Unlifted)), ("t_8324_8329_8320", (56, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8324_8329_8321", (56, TyVarFree NoLoc Unlifted)), ("index_elem_8324_8329_8322", (56, TyVarFree NoLoc Unlifted)), ("kt_8324_8329_8323", (56, TyVarFree NoLoc Lifted)), ("t_8324_8329_8324", (56, TyVarRecord NoLoc (M.fromList [("left", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8324_8329_8323" 15791}) []))]))), ("num_8324_8329_8325", (56, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8325_8320_8320", (56, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8325_8320_8321", (56, TyVarFree NoLoc Unlifted)), ("index_elem_8325_8320_8322", (56, TyVarFree NoLoc Unlifted)), ("t_8325_8320_8323", (57, TyVarFree NoLoc Lifted)), ("t_8325_8320_8324", (58, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("update_elem_8325_8320_8325", (58, TyVarFree NoLoc Unlifted)), ("t_8325_8320_8326", (58, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8325_8320_8327", (58, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8325_8321_8322", (59, TyVarFree NoLoc Lifted)), ("t_8325_8321_8323", (60, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("num_8325_8321_8324", (60, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("update_elem_8325_8321_8325", (60, TyVarFree NoLoc Unlifted)), ("t_8325_8321_8326", (56, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8325_8321_8327", (56, TyVarFree NoLoc Unlifted)), ("index_elem_8325_8321_8328", (56, TyVarFree NoLoc Unlifted)), ("kt_8325_8321_8329", (56, TyVarFree NoLoc Lifted)), ("t_8325_8322_8320", (56, TyVarRecord NoLoc (M.fromList [("left", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8325_8321_8329" 15820}) []))]))), ("t_8325_8322_8323", (57, TyVarFree NoLoc Lifted)), ("t_8325_8322_8324", (58, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8325_8322_8325", (58, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8325_8323_8320", (59, TyVarFree NoLoc Lifted)), ("t_8325_8323_8321", (60, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("num_8325_8323_8322", (60, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("update_elem_8325_8323_8323", (60, TyVarFree NoLoc Unlifted)), ("if_t_8325_8323_8324", (56, TyVarFree NoLoc SizeLifted)), ("t_8325_8323_8325", (56, TyVarFree NoLoc Unlifted)), ("t_8325_8323_8326", (56, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8325_8323_8327", (56, TyVarFree NoLoc Unlifted)), ("index_elem_8325_8323_8328", (56, TyVarFree NoLoc Unlifted)), ("num_8325_8323_8329", (56, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8325_8324_8324", (56, TyVarFree NoLoc Unlifted)), ("t_8325_8324_8325", (56, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8325_8324_8326", (56, TyVarFree NoLoc Unlifted)), ("index_elem_8325_8324_8327", (56, TyVarFree NoLoc Unlifted)), ("kt_8325_8324_8328", (56, TyVarFree NoLoc Lifted)), ("t_8325_8324_8329", (56, TyVarRecord NoLoc (M.fromList [("right", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8325_8324_8328" 15852}) []))]))), ("num_8325_8325_8320", (56, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8325_8325_8325", (56, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8325_8325_8326", (56, TyVarFree NoLoc Unlifted)), ("index_elem_8325_8325_8327", (56, TyVarFree NoLoc Unlifted)), ("t_8325_8325_8328", (57, TyVarFree NoLoc Lifted)), ("t_8325_8325_8329", (58, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("t_8325_8326_8320", (58, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8325_8326_8321", (58, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8325_8326_8322", (58, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8325_8326_8323", (58, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8325_8326_8324", (58, TyVarFree NoLoc Unlifted)), ("index_elem_8325_8326_8325", (58, TyVarFree NoLoc Unlifted)), ("kt_8325_8326_8326", (58, TyVarFree NoLoc Lifted)), ("t_8325_8326_8327", (58, TyVarRecord NoLoc (M.fromList [("left", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8325_8326_8326" 15871}) []))]))), ("num_8325_8326_8328", (58, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8325_8327_8325", (58, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("update_elem_8325_8328_8324", (58, TyVarFree NoLoc Unlifted)), ("t_8325_8328_8325", (58, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8325_8328_8326", (58, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8325_8329_8321", (59, TyVarFree NoLoc Lifted)), ("t_8325_8329_8322", (60, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("num_8325_8329_8323", (60, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("update_elem_8325_8329_8324", (60, TyVarFree NoLoc Unlifted)), ("t_8325_8329_8325", (56, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8325_8329_8326", (56, TyVarFree NoLoc Unlifted)), ("index_elem_8325_8329_8327", (56, TyVarFree NoLoc Unlifted)), ("kt_8325_8329_8328", (56, TyVarFree NoLoc Lifted)), ("t_8325_8329_8329", (56, TyVarRecord NoLoc (M.fromList [("right", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8325_8329_8328" 15909}) []))]))), ("t_8326_8320_8322", (57, TyVarFree NoLoc Lifted)), ("t_8326_8320_8323", (58, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8326_8320_8324", (58, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8326_8320_8329", (59, TyVarFree NoLoc Lifted)), ("t_8326_8321_8320", (60, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("num_8326_8321_8321", (60, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("update_elem_8326_8321_8322", (60, TyVarFree NoLoc Unlifted)), ("if_t_8326_8321_8323", (56, TyVarFree NoLoc SizeLifted)), ("t_8326_8321_8324", (56, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8326_8321_8325", (56, TyVarFree NoLoc Unlifted)), ("index_elem_8326_8321_8326", (56, TyVarFree NoLoc Unlifted)), ("kt_8326_8321_8327", (56, TyVarFree NoLoc Lifted)), ("t_8326_8321_8328", (56, TyVarRecord NoLoc (M.fromList [("parent", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8326_8321_8327" 15931}) []))]))), ("t_8326_8321_8329", (57, TyVarFree NoLoc Lifted)), ("t_8326_8322_8320", (58, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8326_8322_8321", (58, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8326_8322_8326", (59, TyVarFree NoLoc Lifted)), ("t_8326_8322_8327", (60, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("t_8326_8322_8328", (60, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8326_8322_8329", (60, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8326_8323_8320", (60, TyVarFree NoLoc Unlifted)), ("index_elem_8326_8323_8321", (60, TyVarFree NoLoc Unlifted)), ("num_8326_8323_8322", (60, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("update_elem_8326_8323_8327", (60, TyVarFree NoLoc Unlifted)), ("if_t_8326_8323_8328", (56, TyVarFree NoLoc SizeLifted)), ("if_t_8326_8323_8329", (56, TyVarFree NoLoc SizeLifted)), ("t_8326_8324_8320", (57, TyVarFree NoLoc Lifted)), ("t_8326_8324_8321", (57, TyVarFree NoLoc Lifted)), ("t_8326_8324_8322", (57, TyVarFree NoLoc Lifted)), ("t_8326_8324_8323", (57, TyVarFree NoLoc Lifted)), ("t_8326_8324_8324", (57, TyVarFree NoLoc Lifted)), ("if_t_8326_8324_8325", (48, TyVarFree NoLoc SizeLifted)), ("t_8326_8324_8326", (47, TyVarFree NoLoc Lifted)), ("t_8326_8325_8325", (18, TyVarFree NoLoc Lifted)), ("t_8326_8325_8325_8326_8325_8327", (18, TyVarFree NoLoc Lifted)), ("t_8322_8320_8326_8326_8325_8328", (16, TyVarFree NoLoc Lifted)), ("t_8326_8324_8326_8326_8325_8329", (47, TyVarFree NoLoc Lifted)), ("if_t_8326_8324_8325_8326_8326_8320", (48, TyVarFree NoLoc SizeLifted)), ("t_8324_8324_8326_8326_8326_8321", (47, TyVarFree NoLoc Lifted)), ("t_8326_8324_8320_8326_8326_8322", (57, TyVarFree NoLoc Lifted)), ("t_8324_8326_8325_8326_8326_8323", (56, TyVarFree NoLoc Lifted)), ("t_8324_8326_8326_8326_8326_8324", (56, TyVarFree NoLoc Lifted)), ("t_8326_8324_8321_8326_8326_8325", (57, TyVarFree NoLoc Lifted)), ("t_8324_8326_8323_8326_8326_8326", (53, TyVarFree NoLoc Lifted)), ("t_8327_8322_8326_8326_8327", (6, TyVarFree NoLoc Lifted)), ("t_8321_8325_8326_8326_8326_8328", (12, TyVarFree NoLoc Lifted)), ("b_8321_8328_8321_8326_8326_8329", (15, TyVarFree NoLoc Lifted)), ("a_8321_8328_8320_8326_8327_8320", (15, TyVarFree NoLoc Lifted)), ("t_8321_8327_8329_8326_8327_8321", (14, TyVarFree NoLoc Lifted)), ("t_8321_8327_8328_8326_8327_8322", (14, TyVarFree NoLoc Lifted)), ("a_8321_8325_8327_8326_8327_8323", (13, TyVarFree NoLoc Lifted)), ("t_8321_8321_8329_8326_8327_8324", (8, TyVarFree NoLoc Lifted)), ("b_8327_8324_8326_8327_8325", (7, TyVarFree NoLoc Lifted)), ("a_8327_8323_8326_8327_8326", (7, TyVarFree NoLoc Lifted)), ("b_8327_8326_8326_8327_8327", (7, TyVarFree NoLoc Lifted)), ("a_8327_8325_8326_8327_8328", (7, TyVarFree NoLoc Lifted))]+    ),+    ( [ "t_8320" ~ "t_8323",+        "num_8324" ~ "t_8323",+        "t_8323" ~ "f32",+        "num_8329" ~ "f32",+        "f32" ~ "f32",+        "num_8321_8324" ~ "f32",+        "t_8321_8329" ~ "f32",+        "t_8321" ~ "t_8322_8320",+        "num_8322_8321" ~ "t_8322_8320",+        "t_8322_8320" ~ "f32",+        "num_8322_8326" ~ "f32",+        "f32" ~ "f32",+        "num_8323_8321" ~ "f32",+        "t_8323_8326" ~ "f32",+        "t_8322" ~ "t_8323_8327",+        "num_8323_8328" ~ "t_8323_8327",+        "t_8323_8327" ~ "f32",+        "num_8324_8323" ~ "f32",+        "f32" ~ "f32",+        "num_8324_8328" ~ "f32",+        "t_8325_8323" ~ "f32",+        "t_8321_8329" ~ "f32",+        "u32" ~ "u32",+        "t_8325_8328" ~ "u32",+        "t_8323_8326" ~ "f32",+        "u32" ~ "u32",+        "t_8326_8323" ~ "u32",+        "t_8325_8323" ~ "f32",+        "u32" ~ "u32",+        "t_8326_8328" ~ "u32",+        "t_8325_8328" ~ "t_8327_8321",+        "num_8327_8322" ~ "t_8327_8321",+        "t_8326_8323" ~ "t_8327_8327",+        "num_8327_8328" ~ "t_8327_8327",+        "t_8327_8321" ~ "t_8327_8320",+        "t_8327_8327" ~ "t_8327_8320",+        "t_8327_8320" ~ "t_8326_8329",+        "t_8326_8328" ~ "t_8326_8329",+        "u32" ~ "t_8326_8329"+      ],+      M.empty,+      M.fromList [("t_8320", (1, TyVarFree NoLoc Lifted)), ("t_8321", (1, TyVarFree NoLoc Lifted)), ("t_8322", (1, TyVarFree NoLoc Lifted)), ("t_8323", (2, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8324", (2, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8329", (2, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8324", (2, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8329", (3, TyVarFree NoLoc Lifted)), ("t_8322_8320", (4, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8321", (4, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8326", (4, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8321", (4, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8326", (5, TyVarFree NoLoc Lifted)), ("t_8323_8327", (6, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8328", (6, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8324_8323", (6, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8324_8328", (6, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8325_8323", (7, TyVarFree NoLoc Lifted)), ("t_8325_8328", (9, TyVarFree NoLoc Lifted)), ("t_8326_8323", (11, TyVarFree NoLoc Lifted)), ("t_8326_8328", (13, TyVarFree NoLoc Lifted)), ("t_8326_8329", (14, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8327_8320", (14, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8327_8321", (14, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8327_8322", (14, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8327_8327", (14, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8327_8328", (14, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64]))]+    ),+    ( [ "t_8323_8321_8325" ~ "[]t_8323_8321_8325_8323_8323_8326",+        "t_8321_8325_8327" ~ "[]t_8321_8325_8327_8323_8323_8327",+        "t_8322_8321_8328" ~ "[]t_8322_8321_8328_8323_8323_8328",+        "t_8321_8327_8329" ~ "[]t_8321_8327_8329_8323_8323_8329",+        "kt_8322_8326_8329" ~ "[]kt_8322_8326_8329_8323_8324_8320",+        "kt_8322_8325_8324" ~ "[]kt_8322_8325_8324_8323_8324_8321",+        "range_8322_8322_8327" ~ "[]range_8322_8322_8327_8323_8324_8322",+        "t_8321_8327_8324" ~ "[]t_8321_8327_8324_8323_8324_8323",+        "t_8321_8327_8321" ~ "[]t_8321_8327_8321_8323_8324_8324",+        "ft_8324" ~ "ft_8326",+        "a_8322" ~ "rt_8325",+        "a_8322 -> b_8323" ~ "a_8320 -> x_8321",+        "[]{mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}" ~ "[]a_8320",+        "[]x_8321" ~ "[]f32",+        "t_8321_8323" ~ "f32",+        "ft_8321_8328" ~ "ft_8322_8320",+        "a_8321_8326" ~ "rt_8321_8329",+        "a_8321_8326 -> b_8321_8327" ~ "a_8321_8324 -> x_8321_8325",+        "[]{mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}" ~ "[]a_8321_8324",+        "[]x_8321_8325" ~ "[]f32",+        "t_8322_8327" ~ "f32",+        "ft_8323_8322" ~ "ft_8323_8324",+        "a_8323_8320" ~ "rt_8323_8323",+        "a_8323_8320 -> b_8323_8321" ~ "a_8322_8328 -> x_8322_8329",+        "[]{mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}" ~ "[]a_8322_8328",+        "[]x_8322_8329" ~ "[]f32",+        "t_8324_8321" ~ "f32",+        "ft_8324_8326" ~ "ft_8324_8328",+        "a_8324_8324" ~ "rt_8324_8327",+        "a_8324_8324 -> b_8324_8325" ~ "a_8324_8322 -> x_8324_8323",+        "[]{mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}" ~ "[]a_8324_8322",+        "[]x_8324_8323" ~ "[]f32",+        "t_8325_8325" ~ "f32",+        "ft_8326_8320" ~ "ft_8326_8322",+        "a_8325_8328" ~ "rt_8326_8321",+        "a_8325_8328 -> b_8325_8329" ~ "a_8325_8326 -> x_8325_8327",+        "[]{mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}" ~ "[]a_8325_8326",+        "[]x_8325_8327" ~ "[]f32",+        "t_8326_8329" ~ "f32",+        "ft_8327_8324" ~ "ft_8327_8326",+        "a_8327_8322" ~ "rt_8327_8325",+        "a_8327_8322 -> b_8327_8323" ~ "a_8327_8320 -> x_8327_8321",+        "[]{mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}" ~ "[]a_8327_8320",+        "[]x_8327_8321" ~ "[]f32",+        "t_8328_8323" ~ "f32",+        "t_8328_8324" ~ "t_8328_8328",+        "t_8325_8325" ~ "t_8328_8328",+        "t_8321_8323" ~ "t_8329_8323",+        "t_8325_8325" ~ "t_8329_8323",+        "t_8328_8328" ~ "t_8328_8327",+        "t_8329_8323" ~ "t_8328_8327",+        "t_8328_8325" ~ "t_8321_8320_8323",+        "t_8326_8329" ~ "t_8321_8320_8323",+        "t_8322_8327" ~ "t_8321_8320_8328",+        "t_8326_8329" ~ "t_8321_8320_8328",+        "t_8321_8320_8323" ~ "t_8321_8320_8322",+        "t_8321_8320_8328" ~ "t_8321_8320_8322",+        "t_8328_8326" ~ "t_8321_8321_8328",+        "t_8328_8323" ~ "t_8321_8321_8328",+        "t_8324_8321" ~ "t_8321_8322_8323",+        "t_8328_8323" ~ "t_8321_8322_8323",+        "t_8321_8321_8328" ~ "t_8321_8321_8327",+        "t_8321_8322_8323" ~ "t_8321_8321_8327",+        "{x: t_8328_8324, y: t_8328_8325, z: t_8328_8326} -> {x: t_8328_8327, y: t_8321_8320_8322, z: t_8321_8321_8327}" ~ "a_8321_8323_8322 -> b_8321_8323_8323",+        "{x: f32, y: f32, z: f32} -> u32" ~ "b_8321_8323_8323 -> c_8321_8323_8324",+        "t_8321_8323_8329" ~ "{x: a_8321_8323_8322} -> c_8321_8323_8324",+        "t_8321_8324_8322" ~ "t_8321_8324_8324",+        "t_8321_8323_8329" ~ "arg_8321_8324_8325 -> res_8321_8324_8326",+        "kt_8321_8324_8323" ~ "arg_8321_8324_8325",+        "{p: t_8321_8324_8322} -> res_8321_8324_8326" ~ "t_8321_8324_8320 -> k_8321_8324_8321",+        "i32" ~ "i32",+        "i32 -> u32 -> i32" ~ "i32 -> k_8321_8324_8321 -> i32",+        "[]{mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}" ~ "[]t_8321_8324_8320",+        "t_8321_8325_8327" ~ "[]t_8321_8324_8320",+        "t_8321_8326_8320" ~ "t_8321_8326_8322",+        "t_8321_8323_8329" ~ "arg_8321_8326_8323 -> res_8321_8326_8324",+        "kt_8321_8326_8321" ~ "arg_8321_8326_8323",+        "{p: t_8321_8326_8320} -> res_8321_8326_8324" ~ "a_8321_8325_8328 -> x_8321_8325_8329",+        "t_8321_8325_8327" ~ "[]a_8321_8325_8328",+        "t_8321_8327_8321" ~ "[]x_8321_8325_8329",+        "t_8321_8327_8321" ~ "[]u32",+        "t_8321_8327_8324" ~ "[]{delta_node: i32, left: #inner i32 | #leaf i32, parent: i32, right: #inner i32 | #leaf i32, sfc_code: u32}",+        "t_8321_8325_8327" ~ "[]{mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}",+        "t_8321_8327_8324" ~ "[]{delta_node: i32, left: #inner i32 | #leaf i32, parent: i32, right: #inner i32 | #leaf i32, sfc_code: u32}",+        "t_8321_8327_8329" ~ "[]{body: {mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}, children: []i32, is_leaf: bool, parent: i32, tree_level: i32}",+        "t_8321_8328_8321" ~ "num_8321_8328_8320",+        "index_8321_8328_8322" ~ "index_elem_8321_8328_8323",+        "t_8321_8327_8324" ~ "[]index_elem_8321_8328_8323",+        "index_8321_8328_8322" ~ "t_8321_8328_8325",+        "t_8321_8328_8326" ~ "kt_8321_8328_8324",+        "i32" ~ "t_8321_8328_8328",+        "num_8321_8328_8329" ~ "t_8321_8328_8328",+        "t_8321_8328_8328" ~ "t_8321_8328_8327",+        "t_8321_8328_8326" ~ "t_8321_8328_8327",+        "t_8321_8329_8328" ~ "t_8321_8328_8327",+        "t_8321_8329_8328" ~ "t_8322_8320_8321",+        "num_8322_8320_8322" ~ "t_8322_8320_8321",+        "t_8322_8320_8321" ~ "t_8322_8320_8320",+        "num_8322_8320_8327" ~ "t_8322_8320_8320",+        "t_8322_8320_8320" ~ "t_8321_8329_8329",+        "num_8322_8321_8322" ~ "t_8321_8329_8329",+        "t_8322_8321_8327" ~ "t_8321_8329_8329",+        "t_8322_8321_8328" ~ "t_8321_8327_8329",+        "t_8322_8321_8329" ~ "t_8322_8321_8327",+        "t_8322_8321_8327" ~ "t_8322_8322_8320",+        "num_8322_8322_8321" ~ "t_8322_8322_8320",+        "t_8322_8321_8327" ~ "t_8322_8322_8320",+        "t_8322_8321_8327" ~ "num_8322_8322_8326",+        "range_8322_8322_8327" ~ "[]t_8322_8321_8327",+        "range_8322_8322_8327" ~ "[]elem_8322_8322_8328",+        "t_8322_8322_8329" ~ "elem_8322_8322_8328",+        "{body: {mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}, children: []i32, is_leaf: bool, parent: i32, tree_level: i32}" ~ "t_8322_8323_8324",+        "kt_8322_8323_8323" ~ "t_8322_8323_8322",+        "t_8322_8322_8329" ~ "t_8322_8323_8322",+        "t_8322_8324_8320" ~ "num_8322_8323_8329",+        "{body: {mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}, children: []i32, is_leaf: bool, parent: i32, tree_level: i32}" ~ "t_8322_8324_8322",+        "t_8322_8324_8323" ~ "kt_8322_8324_8321",+        "(t_8322_8324_8324, t_8322_8324_8325)" ~ "(t_8322_8324_8323, t_8322_8324_8320)",+        "t_8322_8324_8325" ~ "t_8322_8324_8327",+        "num_8322_8324_8328" ~ "t_8322_8324_8327",+        "{body: {mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}, children: []i32, is_leaf: bool, parent: i32, tree_level: i32}" ~ "t_8322_8325_8325",+        "t_8322_8325_8326" ~ "t_8322_8324_8325",+        "index_8322_8325_8327" ~ "index_elem_8322_8325_8328",+        "kt_8322_8325_8324" ~ "[]index_elem_8322_8325_8328",+        "t_8322_8326_8320" ~ "num_8322_8325_8329",+        "index_8322_8325_8327" ~ "t_8322_8325_8323",+        "num_8322_8325_8329" ~ "t_8322_8325_8323",+        "bool" ~ "t_8322_8324_8326",+        "bool" ~ "t_8322_8324_8326",+        "{body: {mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}, children: []i32, is_leaf: bool, parent: i32, tree_level: i32}" ~ "t_8322_8327_8320",+        "t_8322_8327_8321" ~ "t_8322_8324_8325",+        "index_8322_8327_8322" ~ "index_elem_8322_8327_8323",+        "kt_8322_8326_8329" ~ "[]index_elem_8322_8327_8323",+        "t_8322_8327_8324" ~ "index_8322_8327_8322",+        "index_8322_8327_8325" ~ "index_elem_8322_8327_8326",+        "t_8322_8321_8328" ~ "[]index_elem_8322_8327_8326",+        "index_8322_8327_8325" ~ "t_8322_8327_8328",+        "t_8322_8327_8329" ~ "kt_8322_8327_8327",+        "t_8322_8324_8324" ~ "t_8322_8328_8321",+        "t_8322_8327_8329" ~ "t_8322_8328_8323",+        "kt_8322_8328_8320" ~ "{x: f32, y: f32, z: f32}",+        "kt_8322_8328_8322" ~ "{x: f32, y: f32, z: f32}",+        "t_8322_8328_8328" ~ "{x: f32, y: f32, z: f32}",+        "t_8322_8324_8324" ~ "t_8322_8329_8321",+        "t_8322_8327_8329" ~ "t_8322_8329_8323",+        "kt_8322_8329_8320" ~ "t_8322_8328_8329",+        "kt_8322_8329_8322" ~ "t_8322_8328_8329",+        "t_8322_8329_8328" ~ "t_8322_8328_8329",+        "t_8322_8324_8324" ~ "t_8323_8320_8320",+        "t_8322_8324_8325" ~ "t_8323_8320_8321",+        "num_8323_8320_8322" ~ "t_8323_8320_8321",+        "(t_8322_8324_8323, t_8322_8324_8320)" ~ "({mass: t_8322_8329_8328, position: t_8322_8328_8328, velocity: kt_8322_8329_8329}, t_8323_8320_8321)",+        "(t_8323_8320_8327, t_8323_8320_8328)" ~ "(t_8322_8324_8324, t_8322_8324_8325)",+        "{body: {mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}, children: []i32, is_leaf: bool, parent: i32, tree_level: i32}" ~ "ft_8323_8320_8329",+        "bool" ~ "bool",+        "{body: {mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}, children: []i32, is_leaf: bool, parent: i32, tree_level: i32}" ~ "if_t_8323_8321_8320",+        "{body: {mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}, children: []i32, is_leaf: bool, parent: i32, tree_level: i32}" ~ "if_t_8323_8321_8320",+        "{n: {body: {mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}, children: []i32, is_leaf: bool, parent: i32, tree_level: i32}} -> if_t_8323_8321_8320" ~ "a_8322_8323_8320 -> x_8322_8323_8321",+        "t_8322_8321_8328" ~ "[]a_8322_8323_8320",+        "t_8321_8327_8329" ~ "[]x_8322_8323_8321",+        "t_8323_8321_8325" ~ "t_8322_8321_8328",+        "t_8321_8323" ~ "t_8323_8321_8327",+        "t_8325_8325" ~ "t_8323_8321_8327",+        "t_8323_8321_8327" ~ "et_8323_8321_8326",+        "t_8322_8327" ~ "t_8323_8322_8322",+        "t_8326_8329" ~ "t_8323_8322_8322",+        "t_8323_8322_8322" ~ "et_8323_8321_8326",+        "t_8324_8321" ~ "t_8323_8322_8327",+        "t_8328_8323" ~ "t_8323_8322_8327",+        "t_8323_8322_8327" ~ "et_8323_8321_8326",+        "[]et_8323_8321_8326" ~ "[]f32",+        "t_8323_8323_8324" ~ "f32",+        "([]{body: {mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}, children: []i32, is_leaf: bool, parent: i32, tree_level: i32}, f32, i32, []{mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32})}" ~ "(t_8323_8321_8325, t_8323_8323_8324, t_8321_8328_8326, t_8321_8325_8327)"+      ],+      M.empty,+      M.fromList [("a_8320", (2, TyVarFree NoLoc Unlifted)), ("x_8321", (2, TyVarFree NoLoc Unlifted)), ("a_8322", (2, TyVarFree NoLoc Lifted)), ("b_8323", (2, TyVarFree NoLoc Lifted)), ("ft_8324", (2, TyVarFree NoLoc Lifted)), ("rt_8325", (2, TyVarRecord NoLoc (M.fromList [("position", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "ft_8324" 16341}) []))]))), ("ft_8326", (2, TyVarRecord NoLoc (M.fromList [("x", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "b_8323" 16340}) []))]))), ("t_8321_8323", (3, TyVarFree NoLoc Lifted)), ("a_8321_8324", (4, TyVarFree NoLoc Unlifted)), ("x_8321_8325", (4, TyVarFree NoLoc Unlifted)), ("a_8321_8326", (4, TyVarFree NoLoc Lifted)), ("b_8321_8327", (4, TyVarFree NoLoc Lifted)), ("ft_8321_8328", (4, TyVarFree NoLoc Lifted)), ("rt_8321_8329", (4, TyVarRecord NoLoc (M.fromList [("position", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "ft_8321_8328" 16361}) []))]))), ("ft_8322_8320", (4, TyVarRecord NoLoc (M.fromList [("y", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "b_8321_8327" 16360}) []))]))), ("t_8322_8327", (5, TyVarFree NoLoc Lifted)), ("a_8322_8328", (6, TyVarFree NoLoc Unlifted)), ("x_8322_8329", (6, TyVarFree NoLoc Unlifted)), ("a_8323_8320", (6, TyVarFree NoLoc Lifted)), ("b_8323_8321", (6, TyVarFree NoLoc Lifted)), ("ft_8323_8322", (6, TyVarFree NoLoc Lifted)), ("rt_8323_8323", (6, TyVarRecord NoLoc (M.fromList [("position", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "ft_8323_8322" 16381}) []))]))), ("ft_8323_8324", (6, TyVarRecord NoLoc (M.fromList [("z", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "b_8323_8321" 16380}) []))]))), ("t_8324_8321", (7, TyVarFree NoLoc Lifted)), ("a_8324_8322", (8, TyVarFree NoLoc Unlifted)), ("x_8324_8323", (8, TyVarFree NoLoc Unlifted)), ("a_8324_8324", (8, TyVarFree NoLoc Lifted)), ("b_8324_8325", (8, TyVarFree NoLoc Lifted)), ("ft_8324_8326", (8, TyVarFree NoLoc Lifted)), ("rt_8324_8327", (8, TyVarRecord NoLoc (M.fromList [("position", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "ft_8324_8326" 16401}) []))]))), ("ft_8324_8328", (8, TyVarRecord NoLoc (M.fromList [("x", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "b_8324_8325" 16400}) []))]))), ("t_8325_8325", (9, TyVarFree NoLoc Lifted)), ("a_8325_8326", (10, TyVarFree NoLoc Unlifted)), ("x_8325_8327", (10, TyVarFree NoLoc Unlifted)), ("a_8325_8328", (10, TyVarFree NoLoc Lifted)), ("b_8325_8329", (10, TyVarFree NoLoc Lifted)), ("ft_8326_8320", (10, TyVarFree NoLoc Lifted)), ("rt_8326_8321", (10, TyVarRecord NoLoc (M.fromList [("position", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "ft_8326_8320" 16421}) []))]))), ("ft_8326_8322", (10, TyVarRecord NoLoc (M.fromList [("y", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "b_8325_8329" 16420}) []))]))), ("t_8326_8329", (11, TyVarFree NoLoc Lifted)), ("a_8327_8320", (12, TyVarFree NoLoc Unlifted)), ("x_8327_8321", (12, TyVarFree NoLoc Unlifted)), ("a_8327_8322", (12, TyVarFree NoLoc Lifted)), ("b_8327_8323", (12, TyVarFree NoLoc Lifted)), ("ft_8327_8324", (12, TyVarFree NoLoc Lifted)), ("rt_8327_8325", (12, TyVarRecord NoLoc (M.fromList [("position", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "ft_8327_8324" 16441}) []))]))), ("ft_8327_8326", (12, TyVarRecord NoLoc (M.fromList [("z", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "b_8327_8323" 16440}) []))]))), ("t_8328_8323", (13, TyVarFree NoLoc Lifted)), ("t_8328_8324", (15, TyVarFree NoLoc Lifted)), ("t_8328_8325", (15, TyVarFree NoLoc Lifted)), ("t_8328_8326", (15, TyVarFree NoLoc Lifted)), ("t_8328_8327", (16, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8328_8328", (16, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8329_8323", (16, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8320_8322", (16, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8320_8323", (16, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8320_8328", (16, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8321_8327", (16, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8321_8328", (16, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8322_8323", (16, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("a_8321_8323_8322", (14, TyVarFree NoLoc Lifted)), ("b_8321_8323_8323", (14, TyVarFree NoLoc Lifted)), ("c_8321_8323_8324", (14, TyVarFree NoLoc Lifted)), ("t_8321_8323_8329", (15, TyVarFree NoLoc Lifted)), ("t_8321_8324_8320", (16, TyVarFree NoLoc Unlifted)), ("k_8321_8324_8321", (16, TyVarFree NoLoc Unlifted)), ("t_8321_8324_8322", (17, TyVarFree NoLoc Lifted)), ("kt_8321_8324_8323", (18, TyVarFree NoLoc Lifted)), ("t_8321_8324_8324", (18, TyVarRecord NoLoc (M.fromList [("position", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8321_8324_8323" 16528}) []))]))), ("arg_8321_8324_8325", (18, TyVarFree NoLoc Lifted)), ("res_8321_8324_8326", (18, TyVarFree NoLoc Lifted)), ("t_8321_8325_8327", (17, TyVarFree NoLoc Lifted)), ("a_8321_8325_8328", (18, TyVarFree NoLoc Unlifted)), ("x_8321_8325_8329", (18, TyVarFree NoLoc Unlifted)), ("t_8321_8326_8320", (19, TyVarFree NoLoc Lifted)), ("kt_8321_8326_8321", (20, TyVarFree NoLoc Lifted)), ("t_8321_8326_8322", (20, TyVarRecord NoLoc (M.fromList [("position", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8321_8326_8321" 16552}) []))]))), ("arg_8321_8326_8323", (20, TyVarFree NoLoc Lifted)), ("res_8321_8326_8324", (20, TyVarFree NoLoc Lifted)), ("t_8321_8327_8321", (19, TyVarFree NoLoc Lifted)), ("t_8321_8327_8324", (21, TyVarFree NoLoc Lifted)), ("t_8321_8327_8329", (23, TyVarFree NoLoc Lifted)), ("num_8321_8328_8320", (24, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8328_8321", (24, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8321_8328_8322", (24, TyVarFree NoLoc Unlifted)), ("index_elem_8321_8328_8323", (24, TyVarFree NoLoc Unlifted)), ("kt_8321_8328_8324", (24, TyVarFree NoLoc Lifted)), ("t_8321_8328_8325", (24, TyVarRecord NoLoc (M.fromList [("delta_node", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8321_8328_8324" 16585}) []))]))), ("t_8321_8328_8326", (25, TyVarFree NoLoc Lifted)), ("t_8321_8328_8327", (26, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8328_8328", (26, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321_8328_8329", (26, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8329_8328", (27, TyVarFree NoLoc Lifted)), ("t_8321_8329_8329", (28, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8320_8320", (28, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8320_8321", (28, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8320_8322", (28, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8320_8327", (28, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8321_8322", (28, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8321_8327", (29, TyVarFree NoLoc Lifted)), ("t_8322_8321_8328", (30, TyVarFree NoLoc Lifted)), ("t_8322_8321_8329", (30, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64])), ("t_8322_8322_8320", (30, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8322_8321", (30, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8322_8326", (30, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("range_8322_8322_8327", (30, TyVarFree NoLoc Unlifted)), ("elem_8322_8322_8328", (30, TyVarFree NoLoc Unlifted)), ("t_8322_8322_8329", (30, TyVarFree NoLoc Lifted)), ("a_8322_8323_8320", (30, TyVarFree NoLoc Unlifted)), ("x_8322_8323_8321", (30, TyVarFree NoLoc Unlifted)), ("t_8322_8323_8322", (32, TyVarFree NoLoc Unlifted)), ("kt_8322_8323_8323", (32, TyVarFree NoLoc Lifted)), ("t_8322_8323_8324", (32, TyVarRecord NoLoc (M.fromList [("tree_level", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8322_8323_8323" 16643}) []))]))), ("num_8322_8323_8329", (32, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8324_8320", (33, TyVarFree NoLoc Lifted)), ("kt_8322_8324_8321", (34, TyVarFree NoLoc Lifted)), ("t_8322_8324_8322", (34, TyVarRecord NoLoc (M.fromList [("body", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8322_8324_8321" 16654}) []))]))), ("t_8322_8324_8323", (35, TyVarFree NoLoc Lifted)), ("t_8322_8324_8324", (36, TyVarFree NoLoc Lifted)), ("t_8322_8324_8325", (36, TyVarFree NoLoc Lifted)), ("t_8322_8324_8326", (36, TyVarPrim NoLoc [Bool])), ("t_8322_8324_8327", (36, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8324_8328", (36, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8325_8323", (36, TyVarFree NoLoc Unlifted)), ("kt_8322_8325_8324", (36, TyVarFree NoLoc Lifted)), ("t_8322_8325_8325", (36, TyVarRecord NoLoc (M.fromList [("children", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8322_8325_8324" 16669}) []))]))), ("t_8322_8325_8326", (36, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8322_8325_8327", (36, TyVarFree NoLoc Unlifted)), ("index_elem_8322_8325_8328", (36, TyVarFree NoLoc Unlifted)), ("num_8322_8325_8329", (36, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8326_8320", (36, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("kt_8322_8326_8329", (36, TyVarFree NoLoc Lifted)), ("t_8322_8327_8320", (36, TyVarRecord NoLoc (M.fromList [("children", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8322_8326_8329" 16687}) []))]))), ("t_8322_8327_8321", (36, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8322_8327_8322", (36, TyVarFree NoLoc Unlifted)), ("index_elem_8322_8327_8323", (36, TyVarFree NoLoc Unlifted)), ("t_8322_8327_8324", (36, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8322_8327_8325", (36, TyVarFree NoLoc Unlifted)), ("index_elem_8322_8327_8326", (36, TyVarFree NoLoc Unlifted)), ("kt_8322_8327_8327", (36, TyVarFree NoLoc Lifted)), ("t_8322_8327_8328", (36, TyVarRecord NoLoc (M.fromList [("body", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8322_8327_8327" 16695}) []))]))), ("t_8322_8327_8329", (37, TyVarFree NoLoc Lifted)), ("kt_8322_8328_8320", (38, TyVarFree NoLoc Lifted)), ("t_8322_8328_8321", (38, TyVarRecord NoLoc (M.fromList [("position", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8322_8328_8320" 16698}) []))]))), ("kt_8322_8328_8322", (38, TyVarFree NoLoc Lifted)), ("t_8322_8328_8323", (38, TyVarRecord NoLoc (M.fromList [("position", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8322_8328_8322" 16700}) []))]))), ("t_8322_8328_8328", (39, TyVarFree NoLoc Lifted)), ("t_8322_8328_8329", (40, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("kt_8322_8329_8320", (40, TyVarFree NoLoc Lifted)), ("t_8322_8329_8321", (40, TyVarRecord NoLoc (M.fromList [("mass", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8322_8329_8320" 16709}) []))]))), ("kt_8322_8329_8322", (40, TyVarFree NoLoc Lifted)), ("t_8322_8329_8323", (40, TyVarRecord NoLoc (M.fromList [("mass", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8322_8329_8322" 16711}) []))]))), ("t_8322_8329_8328", (41, TyVarFree NoLoc Lifted)), ("kt_8322_8329_8329", (42, TyVarFree NoLoc Lifted)), ("t_8323_8320_8320", (42, TyVarRecord NoLoc (M.fromList [("velocity", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8322_8329_8329" 16719}) []))]))), ("t_8323_8320_8321", (42, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8320_8322", (42, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8320_8327", (37, TyVarFree NoLoc Lifted)), ("t_8323_8320_8328", (37, TyVarFree NoLoc Lifted)), ("ft_8323_8320_8329", (38, TyVarRecord NoLoc (M.fromList [("body", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "t_8323_8320_8327" 16731}) []))]))), ("if_t_8323_8321_8320", (32, TyVarFree NoLoc SizeLifted)), ("t_8323_8321_8325", (27, TyVarFree NoLoc Lifted)), ("et_8323_8321_8326", (28, TyVarFree NoLoc Unlifted)), ("t_8323_8321_8327", (28, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8322_8322", (28, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8322_8327", (28, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8323_8324", (29, TyVarFree NoLoc Lifted)), ("t_8323_8321_8325_8323_8323_8326", (27, TyVarFree NoLoc Lifted)), ("t_8321_8325_8327_8323_8323_8327", (17, TyVarFree NoLoc Lifted)), ("t_8322_8321_8328_8323_8323_8328", (30, TyVarFree NoLoc Lifted)), ("t_8321_8327_8329_8323_8323_8329", (23, TyVarFree NoLoc Lifted)), ("kt_8322_8326_8329_8323_8324_8320", (36, TyVarFree NoLoc Lifted)), ("kt_8322_8325_8324_8323_8324_8321", (36, TyVarFree NoLoc Lifted)), ("range_8322_8322_8327_8323_8324_8322", (30, TyVarFree NoLoc Unlifted)), ("t_8321_8327_8324_8323_8324_8323", (21, TyVarFree NoLoc Lifted)), ("t_8321_8327_8321_8323_8324_8324", (19, TyVarFree NoLoc Lifted))]+    ),+    ( [ "t_8322_8326_8324" ~ "[]t_8322_8326_8324_8323_8322_8320",+        "a_8322_8326_8325" ~ "[]a_8322_8326_8325_8323_8322_8321",+        "kt_8322_8326_8322" ~ "[]kt_8322_8326_8322_8323_8322_8322",+        "kt_8329_8321" ~ "[]kt_8329_8321_8323_8322_8323",+        "kt_8326_8328" ~ "[]kt_8326_8328_8323_8322_8324",+        "kt_8324_8324" ~ "[]kt_8324_8324_8323_8322_8325",+        "t_8322" ~ "num_8321",+        "(t_8324, t_8325, t_8326, t_8327)" ~ "({x: f32, y: f32, z: f32}, num_8320, num_8321, num_8323)",+        "t_8327" ~ "t_8328",+        "num_8329" ~ "t_8328",+        "t_8321_8327" ~ "num_8321_8326",+        "t_8326" ~ "t_8321_8325",+        "num_8321_8326" ~ "t_8321_8325",+        "t_8322_8323" ~ "t_8326",+        "index_8322_8324" ~ "index_elem_8322_8325",+        "[]{body: {mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}, children: []i32, is_leaf: bool, parent: i32, tree_level: i32}" ~ "[]index_elem_8322_8325",+        "index_8322_8324" ~ "t_8322_8327",+        "kt_8322_8326" ~ "t_8322_8322",+        "t_8325" ~ "t_8322_8322",+        "bool" ~ "t_8321_8324",+        "bool" ~ "t_8321_8324",+        "t_8323_8326" ~ "t_8321_8324",+        "t_8323_8327" ~ "t_8325",+        "index_8323_8328" ~ "index_elem_8323_8329",+        "[]{body: {mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}, children: []i32, is_leaf: bool, parent: i32, tree_level: i32}" ~ "[]index_elem_8323_8329",+        "t_8324_8320" ~ "index_8323_8328",+        "t_8324_8322" ~ "num_8324_8321",+        "t_8324_8320" ~ "t_8324_8325",+        "t_8324_8326" ~ "t_8324_8322",+        "index_8324_8327" ~ "index_elem_8324_8328",+        "kt_8324_8324" ~ "[]index_elem_8324_8328",+        "index_8324_8327" ~ "t_8324_8323",+        "t_8326" ~ "t_8324_8323",+        "t_8324_8322" ~ "t_8325_8323",+        "num_8325_8324" ~ "t_8325_8323",+        "num_8324_8321" ~ "t_8325_8323",+        "t_8325_8329" ~ "t_8324_8322",+        "t_8325_8329" ~ "t_8326_8321",+        "num_8326_8322" ~ "t_8326_8321",+        "t_8324_8320" ~ "t_8326_8329",+        "t_8325_8329" ~ "t_8327_8320",+        "num_8327_8321" ~ "t_8327_8320",+        "t_8327_8326" ~ "t_8327_8320",+        "index_8327_8327" ~ "index_elem_8327_8328",+        "kt_8326_8328" ~ "[]index_elem_8327_8328",+        "t_8328_8320" ~ "num_8327_8329",+        "index_8327_8327" ~ "t_8326_8327",+        "num_8327_8329" ~ "t_8326_8327",+        "bool" ~ "t_8326_8320",+        "bool" ~ "t_8326_8320",+        "t_8324_8320" ~ "t_8329_8320",+        "t_8324_8320" ~ "t_8329_8322",+        "t_8325_8329" ~ "t_8329_8323",+        "num_8329_8324" ~ "t_8329_8323",+        "t_8329_8329" ~ "t_8329_8323",+        "index_8321_8320_8320" ~ "index_elem_8321_8320_8321",+        "kt_8329_8321" ~ "[]index_elem_8321_8320_8321",+        "t_8326_8320" ~ "bool",+        "kt_8328_8329" ~ "if_t_8321_8320_8322",+        "index_8321_8320_8320" ~ "if_t_8321_8320_8322",+        "t_8321_8320_8323" ~ "if_t_8321_8320_8322",+        "t_8324_8320" ~ "t_8321_8320_8327",+        "i32" ~ "t_8321_8320_8325",+        "kt_8321_8320_8326" ~ "t_8321_8320_8325",+        "t_8321_8320_8325" ~ "i32",+        "f32" ~ "t_8321_8320_8324",+        "f32" ~ "t_8321_8320_8324",+        "t_8321_8321_8328" ~ "t_8321_8320_8324",+        "t_8324_8320" ~ "t_8321_8322_8320",+        "kt_8321_8321_8329" ~ "t_8321_8322_8322",+        "t_8321_8322_8323" ~ "kt_8321_8322_8321",+        "t_8324_8320" ~ "t_8321_8322_8325",+        "kt_8321_8322_8324" ~ "t_8321_8322_8327",+        "t_8321_8322_8328" ~ "kt_8321_8322_8326",+        "f32" ~ "t_8321_8322_8329",+        "t_8321_8322_8328" ~ "t_8321_8322_8329",+        "t_8321_8322_8329" ~ "f32",+        "t_8321_8322_8323" ~ "{x: f32, y: f32, z: f32}",+        "t_8321_8323_8328" ~ "{x: f32, y: f32, z: f32}",+        "{mass: f32, position: {x: f32, y: f32, z: f32}}" ~ "t_8321_8324_8320",+        "t_8321_8324_8321" ~ "kt_8321_8323_8329",+        "t_8321_8324_8321" ~ "t_8321_8324_8324",+        "t_8321_8323_8328" ~ "t_8321_8324_8326",+        "kt_8321_8324_8323" ~ "t_8321_8324_8322",+        "kt_8321_8324_8325" ~ "t_8321_8324_8322",+        "t_8321_8325_8321" ~ "t_8321_8324_8322",+        "t_8321_8324_8321" ~ "t_8321_8325_8324",+        "t_8321_8323_8328" ~ "t_8321_8325_8326",+        "kt_8321_8325_8323" ~ "t_8321_8325_8322",+        "kt_8321_8325_8325" ~ "t_8321_8325_8322",+        "t_8321_8326_8321" ~ "t_8321_8325_8322",+        "t_8321_8324_8321" ~ "t_8321_8326_8324",+        "t_8321_8323_8328" ~ "t_8321_8326_8326",+        "kt_8321_8326_8323" ~ "t_8321_8326_8322",+        "kt_8321_8326_8325" ~ "t_8321_8326_8322",+        "t_8321_8327_8321" ~ "t_8321_8326_8322",+        "t_8321_8325_8321" ~ "t_8321_8327_8324",+        "t_8321_8325_8321" ~ "t_8321_8327_8324",+        "t_8321_8326_8321" ~ "t_8321_8327_8329",+        "t_8321_8326_8321" ~ "t_8321_8327_8329",+        "t_8321_8327_8324" ~ "t_8321_8327_8323",+        "t_8321_8327_8329" ~ "t_8321_8327_8323",+        "t_8321_8327_8321" ~ "t_8321_8328_8328",+        "t_8321_8327_8321" ~ "t_8321_8328_8328",+        "t_8321_8327_8323" ~ "t_8321_8327_8322",+        "t_8321_8328_8328" ~ "t_8321_8327_8322",+        "t_8321_8329_8327" ~ "t_8321_8327_8322",+        "t_8321_8329_8327" ~ "f32",+        "t_8321_8321_8328" ~ "t_8321_8329_8328",+        "f32" ~ "t_8321_8329_8328",+        "t_8322_8320_8325" ~ "t_8321_8329_8328",+        "i32" ~ "t_8322_8320_8328",+        "num_8322_8320_8329" ~ "t_8322_8320_8328",+        "t_8324_8320" ~ "t_8322_8321_8325",+        "t_8322_8320_8328" ~ "t_8322_8320_8327",+        "kt_8322_8321_8324" ~ "t_8322_8320_8327",+        "t_8322_8320_8327" ~ "t_8322_8320_8326",+        "num_8322_8322_8320" ~ "t_8322_8320_8326",+        "t_8322_8322_8325" ~ "bool",+        "t_8324_8320" ~ "t_8322_8322_8329",+        "kt_8322_8322_8328" ~ "t_8322_8322_8327",+        "t_8322_8322_8325" ~ "t_8322_8322_8327",+        "t_8322_8320_8325" ~ "t_8322_8323_8324",+        "f32" ~ "t_8322_8323_8324",+        "t_8322_8322_8327" ~ "t_8322_8322_8326",+        "bool" ~ "t_8322_8322_8326",+        "f32" ~ "f32",+        "{mass: f32, position: {x: f32, y: f32, z: f32}}" ~ "{mass: f32, position: {x: f32, y: f32, z: f32}}",+        "{mass: t_8321_8322_8328, position: t_8321_8323_8328}" ~ "{mass: f32, position: {x: f32, y: f32, z: f32}}",+        "t_8322_8324_8329" ~ "{x: f32, y: f32, z: f32}",+        "t_8324" ~ "{x: f32, y: f32, z: f32}",+        "t_8322_8324_8329" ~ "{x: f32, y: f32, z: f32}",+        "t_8324_8320" ~ "t_8322_8325_8325",+        "t_8327" ~ "t_8322_8325_8326",+        "num_8322_8325_8327" ~ "t_8322_8325_8326",+        "t_8324_8320" ~ "t_8322_8326_8323",+        "t_8322_8326_8324" ~ "kt_8322_8326_8322",+        "t_8322_8327_8324" ~ "t_8322_8327_8322",+        "num_8322_8327_8323" ~ "t_8322_8327_8322",+        "bool -> i32" ~ "b_8322_8327_8320 -> c_8322_8327_8321",+        "t_8322_8327_8324 -> bool" ~ "a_8322_8326_8329 -> b_8322_8327_8320",+        "{x: a_8322_8326_8329} -> c_8322_8327_8321" ~ "a_8322_8326_8327 -> x_8322_8326_8328",+        "t_8322_8326_8324" ~ "[]a_8322_8326_8327",+        "t_8322_8328_8328 -> t_8322_8328_8328 -> t_8322_8328_8328" ~ "a_8322_8328_8327 -> a_8322_8328_8327 -> a_8322_8328_8327",+        "num_8322_8328_8329" ~ "a_8322_8328_8327",+        "[]x_8322_8326_8328" ~ "a_8322_8326_8325",+        "{as: []a_8322_8328_8327} -> a_8322_8328_8327" ~ "a_8322_8326_8325 -> b_8322_8326_8326",+        "t_8322_8329_8328" ~ "b_8322_8326_8326",+        "t_8323_8320_8320" ~ "num_8322_8329_8329",+        "index_8323_8320_8321" ~ "index_elem_8323_8320_8322",+        "t_8322_8326_8324" ~ "[]index_elem_8323_8320_8322",+        "t_8327" ~ "t_8323_8320_8324",+        "t_8322_8329_8328" ~ "t_8323_8320_8324",+        "t_8323_8320_8324" ~ "t_8323_8320_8323",+        "num_8323_8320_8329" ~ "t_8323_8320_8323",+        "t_8322_8322_8326" ~ "bool",+        "({x: f32, y: f32, z: f32}, kt_8322_8325_8324, t_8325, t_8322_8325_8326)" ~ "if_t_8323_8321_8324",+        "(t_8324, index_8323_8320_8321, t_8325, t_8323_8320_8323)" ~ "if_t_8323_8321_8324",+        "t_8323_8326" ~ "bool",+        "(t_8324, t_8321_8320_8323, t_8325, t_8327)" ~ "if_t_8323_8321_8325",+        "if_t_8323_8321_8324" ~ "if_t_8323_8321_8325",+        "({x: f32, y: f32, z: f32}, num_8320, num_8321, num_8323)" ~ "if_t_8323_8321_8325",+        "(t_8323_8321_8326, t_8323_8321_8327, t_8323_8321_8328, t_8323_8321_8329)" ~ "(t_8324, t_8325, t_8326, t_8327)",+        "{x: f32, y: f32, z: f32}" ~ "t_8323_8321_8326"+      ],+      M.empty,+      M.fromList [("num_8320", (5, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8321", (5, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322", (5, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323", (5, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8324", (5, TyVarFree NoLoc Lifted)), ("t_8325", (5, TyVarFree NoLoc Lifted)), ("t_8326", (5, TyVarFree NoLoc Lifted)), ("t_8327", (5, TyVarFree NoLoc Lifted)), ("t_8328", (5, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8329", (5, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8324", (5, TyVarPrim NoLoc [Bool])), ("t_8321_8325", (5, TyVarFree NoLoc Unlifted)), ("num_8321_8326", (5, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8327", (5, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8322", (5, TyVarFree NoLoc Unlifted)), ("t_8322_8323", (5, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8322_8324", (5, TyVarFree NoLoc Unlifted)), ("index_elem_8322_8325", (5, TyVarFree NoLoc Unlifted)), ("kt_8322_8326", (5, TyVarFree NoLoc Lifted)), ("t_8322_8327", (5, TyVarRecord NoLoc (M.fromList [("parent", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8322_8326" 17047}) []))]))), ("t_8323_8326", (6, TyVarFree NoLoc Lifted)), ("t_8323_8327", (7, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8323_8328", (7, TyVarFree NoLoc Unlifted)), ("index_elem_8323_8329", (7, TyVarFree NoLoc Unlifted)), ("t_8324_8320", (8, TyVarFree NoLoc Lifted)), ("num_8324_8321", (9, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8324_8322", (9, TyVarFree NoLoc Lifted)), ("t_8324_8323", (9, TyVarFree NoLoc Unlifted)), ("kt_8324_8324", (9, TyVarFree NoLoc Lifted)), ("t_8324_8325", (9, TyVarRecord NoLoc (M.fromList [("children", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8324_8324" 17068}) []))]))), ("t_8324_8326", (9, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8324_8327", (9, TyVarFree NoLoc Unlifted)), ("index_elem_8324_8328", (9, TyVarFree NoLoc Unlifted)), ("t_8325_8323", (9, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8325_8324", (9, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8325_8329", (10, TyVarFree NoLoc Lifted)), ("t_8326_8320", (11, TyVarPrim NoLoc [Bool])), ("t_8326_8321", (11, TyVarFree NoLoc Unlifted)), ("num_8326_8322", (11, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8326_8327", (11, TyVarFree NoLoc Unlifted)), ("kt_8326_8328", (11, TyVarFree NoLoc Lifted)), ("t_8326_8329", (11, TyVarRecord NoLoc (M.fromList [("children", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8326_8328" 17095}) []))]))), ("t_8327_8320", (11, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8327_8321", (11, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8327_8326", (11, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8327_8327", (11, TyVarFree NoLoc Unlifted)), ("index_elem_8327_8328", (11, TyVarFree NoLoc Unlifted)), ("num_8327_8329", (11, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8328_8320", (11, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("kt_8328_8329", (11, TyVarFree NoLoc Lifted)), ("t_8329_8320", (11, TyVarRecord NoLoc (M.fromList [("parent", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8328_8329" 17119}) []))]))), ("kt_8329_8321", (11, TyVarFree NoLoc Lifted)), ("t_8329_8322", (11, TyVarRecord NoLoc (M.fromList [("children", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8329_8321" 17121}) []))]))), ("t_8329_8323", (11, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8329_8324", (11, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8329_8329", (11, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8321_8320_8320", (11, TyVarFree NoLoc Unlifted)), ("index_elem_8321_8320_8321", (11, TyVarFree NoLoc Unlifted)), ("if_t_8321_8320_8322", (11, TyVarFree NoLoc SizeLifted)), ("t_8321_8320_8323", (12, TyVarFree NoLoc Lifted)), ("t_8321_8320_8324", (9, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8320_8325", (9, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("kt_8321_8320_8326", (9, TyVarFree NoLoc Lifted)), ("t_8321_8320_8327", (9, TyVarRecord NoLoc (M.fromList [("tree_level", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8321_8320_8326" 17137}) []))]))), ("t_8321_8321_8328", (10, TyVarFree NoLoc Lifted)), ("kt_8321_8321_8329", (11, TyVarFree NoLoc Lifted)), ("t_8321_8322_8320", (11, TyVarRecord NoLoc (M.fromList [("body", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8321_8321_8329" 17153}) []))]))), ("kt_8321_8322_8321", (11, TyVarFree NoLoc Lifted)), ("t_8321_8322_8322", (11, TyVarRecord NoLoc (M.fromList [("position", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8321_8322_8321" 17155}) []))]))), ("t_8321_8322_8323", (12, TyVarFree NoLoc Lifted)), ("kt_8321_8322_8324", (13, TyVarFree NoLoc Lifted)), ("t_8321_8322_8325", (13, TyVarRecord NoLoc (M.fromList [("body", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8321_8322_8324" 17158}) []))]))), ("kt_8321_8322_8326", (13, TyVarFree NoLoc Lifted)), ("t_8321_8322_8327", (13, TyVarRecord NoLoc (M.fromList [("mass", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8321_8322_8326" 17160}) []))]))), ("t_8321_8322_8328", (14, TyVarFree NoLoc Lifted)), ("t_8321_8322_8329", (15, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8323_8328", (16, TyVarFree NoLoc Lifted)), ("kt_8321_8323_8329", (17, TyVarFree NoLoc Lifted)), ("t_8321_8324_8320", (17, TyVarRecord NoLoc (M.fromList [("position", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8321_8323_8329" 17175}) []))]))), ("t_8321_8324_8321", (18, TyVarFree NoLoc Lifted)), ("t_8321_8324_8322", (19, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("kt_8321_8324_8323", (19, TyVarFree NoLoc Lifted)), ("t_8321_8324_8324", (19, TyVarRecord NoLoc (M.fromList [("x", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8321_8324_8323" 17179}) []))]))), ("kt_8321_8324_8325", (19, TyVarFree NoLoc Lifted)), ("t_8321_8324_8326", (19, TyVarRecord NoLoc (M.fromList [("x", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8321_8324_8325" 17181}) []))]))), ("t_8321_8325_8321", (20, TyVarFree NoLoc Lifted)), ("t_8321_8325_8322", (21, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("kt_8321_8325_8323", (21, TyVarFree NoLoc Lifted)), ("t_8321_8325_8324", (21, TyVarRecord NoLoc (M.fromList [("y", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8321_8325_8323" 17190}) []))]))), ("kt_8321_8325_8325", (21, TyVarFree NoLoc Lifted)), ("t_8321_8325_8326", (21, TyVarRecord NoLoc (M.fromList [("y", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8321_8325_8325" 17192}) []))]))), ("t_8321_8326_8321", (22, TyVarFree NoLoc Lifted)), ("t_8321_8326_8322", (23, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("kt_8321_8326_8323", (23, TyVarFree NoLoc Lifted)), ("t_8321_8326_8324", (23, TyVarRecord NoLoc (M.fromList [("z", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8321_8326_8323" 17201}) []))]))), ("kt_8321_8326_8325", (23, TyVarFree NoLoc Lifted)), ("t_8321_8326_8326", (23, TyVarRecord NoLoc (M.fromList [("z", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8321_8326_8325" 17203}) []))]))), ("t_8321_8327_8321", (24, TyVarFree NoLoc Lifted)), ("t_8321_8327_8322", (25, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8327_8323", (25, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8327_8324", (25, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8327_8329", (25, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8328_8328", (25, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321_8329_8327", (20, TyVarFree NoLoc Lifted)), ("t_8321_8329_8328", (21, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8320_8325", (22, TyVarFree NoLoc Lifted)), ("t_8322_8320_8326", (23, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8320_8327", (23, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8320_8328", (23, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8320_8329", (23, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("kt_8322_8321_8324", (23, TyVarFree NoLoc Lifted)), ("t_8322_8321_8325", (23, TyVarRecord NoLoc (M.fromList [("tree_level", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8322_8321_8324" 17264}) []))]))), ("num_8322_8322_8320", (23, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8322_8325", (24, TyVarFree NoLoc Lifted)), ("t_8322_8322_8326", (25, TyVarPrim NoLoc [Bool])), ("t_8322_8322_8327", (25, TyVarPrim NoLoc [Bool])), ("kt_8322_8322_8328", (25, TyVarFree NoLoc Lifted)), ("t_8322_8322_8329", (25, TyVarRecord NoLoc (M.fromList [("is_leaf", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8322_8322_8328" 17280}) []))]))), ("t_8322_8323_8324", (25, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8324_8329", (26, TyVarFree NoLoc Lifted)), ("kt_8322_8325_8324", (27, TyVarFree NoLoc Lifted)), ("t_8322_8325_8325", (27, TyVarRecord NoLoc (M.fromList [("parent", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8322_8325_8324" 17311}) []))]))), ("t_8322_8325_8326", (27, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8325_8327", (27, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("kt_8322_8326_8322", (25, TyVarFree NoLoc Lifted)), ("t_8322_8326_8323", (25, TyVarRecord NoLoc (M.fromList [("children", Scalar (TypeVar NoUniqueness (QualName {qualQuals = [], qualLeaf = VName "kt_8322_8326_8322" 17322}) []))]))), ("t_8322_8326_8324", (26, TyVarFree NoLoc Lifted)), ("a_8322_8326_8325", (27, TyVarFree NoLoc Lifted)), ("b_8322_8326_8326", (27, TyVarFree NoLoc Lifted)), ("a_8322_8326_8327", (27, TyVarFree NoLoc Unlifted)), ("x_8322_8326_8328", (27, TyVarFree NoLoc Unlifted)), ("a_8322_8326_8329", (27, TyVarFree NoLoc Lifted)), ("b_8322_8327_8320", (27, TyVarFree NoLoc Lifted)), ("c_8322_8327_8321", (27, TyVarFree NoLoc Lifted)), ("t_8322_8327_8322", (27, TyVarPrim NoLoc [Bool, Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8327_8323", (27, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8327_8324", (27, TyVarFree NoLoc Lifted)), ("a_8322_8328_8327", (27, TyVarFree NoLoc Unlifted)), ("t_8322_8328_8328", (27, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8322_8328_8329", (27, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8322_8329_8328", (28, TyVarFree NoLoc Lifted)), ("num_8322_8329_8329", (29, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8320_8320", (29, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])), ("index_8323_8320_8321", (29, TyVarFree NoLoc Unlifted)), ("index_elem_8323_8320_8322", (29, TyVarFree NoLoc Unlifted)), ("t_8323_8320_8323", (29, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8323_8320_8324", (29, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("num_8323_8320_8329", (29, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])), ("if_t_8323_8321_8324", (25, TyVarFree NoLoc SizeLifted)), ("if_t_8323_8321_8325", (9, TyVarFree NoLoc SizeLifted)), ("t_8323_8321_8326", (6, TyVarFree NoLoc Lifted)), ("t_8323_8321_8327", (6, TyVarFree NoLoc Lifted)), ("t_8323_8321_8328", (6, TyVarFree NoLoc Lifted)), ("t_8323_8321_8329", (6, TyVarFree NoLoc Lifted)), ("t_8322_8326_8324_8323_8322_8320", (26, TyVarFree NoLoc Lifted)), ("a_8322_8326_8325_8323_8322_8321", (27, TyVarFree NoLoc Lifted)), ("kt_8322_8326_8322_8323_8322_8322", (25, TyVarFree NoLoc Lifted)), ("kt_8329_8321_8323_8322_8323", (11, TyVarFree NoLoc Lifted)), ("kt_8326_8328_8323_8322_8324", (11, TyVarFree NoLoc Lifted)), ("kt_8324_8324_8323_8322_8325", (9, TyVarFree NoLoc Lifted))]+    ),+    ( [ "t_8326_8327" ~ "[]t_8326_8327_8327_8328",+        "t_8326_8328" ~ "[]t_8326_8328_8327_8329",+        "t_8326_8329" ~ "[]t_8326_8329_8328_8320",+        "t_8326_8320" ~ "[]t_8326_8320_8328_8321",+        "t_8327_8325" ~ "[]t_8327_8325_8328_8322",+        "t_8327_8326" ~ "[]t_8327_8326_8328_8323",+        "t_8327_8327" ~ "[]t_8327_8327_8328_8324",+        "t_8326_8321" ~ "[]t_8326_8321_8328_8325",+        "t_8325_8329" ~ "[]t_8325_8329_8328_8326",+        "a_8324_8324" ~ "[]a_8324_8324_8328_8327",+        "t_8324_8323" ~ "[]t_8324_8323_8328_8328",+        "t_8323_8322" ~ "[]t_8323_8322_8328_8329",+        "t_8321" ~ "float_8320",+        "[]f32" ~ "[]a_8326",+        "[]f32" ~ "[]b_8327",+        "[]f32" ~ "[]c_8328",+        "[]f32" ~ "[]a_8321_8325",+        "[]f32" ~ "[]b_8321_8326",+        "[]f32" ~ "[]c_8321_8327",+        "(f32, f32, f32) -> {mass: f32} -> (f32, f32, f32) -> {mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}" ~ "a_8322 -> b_8323 -> c_8324 -> x_8325",+        "[](a_8326, b_8327, c_8328)" ~ "[]a_8322",+        "[]f32" ~ "[]b_8323",+        "[](a_8321_8325, b_8321_8326, c_8321_8327)" ~ "[]c_8324",+        "t_8323_8322" ~ "[]x_8325",+        "i32" ~ "i32",+        "t_8321" ~ "f32",+        "f32" ~ "f32",+        "f32" ~ "f32",+        "t_8323_8322" ~ "[]{mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}",+        "t_8324_8323" ~ "[]{mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}",+        "{b: {mass: f32, position: {x: f32, y: f32, z: f32}, velocity: {x: f32, y: f32, z: f32}}} -> ((f32, f32, f32), f32, (f32, f32, f32))" ~ "a_8324_8326 -> x_8324_8327",+        "t_8324_8323" ~ "[]a_8324_8326",+        "[]x_8324_8327" ~ "a_8324_8324",+        "{xs: [](a_8325_8322, b_8325_8323, c_8325_8324)} -> ([]a_8325_8322, []b_8325_8323, []c_8325_8324)" ~ "a_8324_8324 -> b_8324_8325",+        "(t_8325_8329, t_8326_8320, t_8326_8321)" ~ "b_8324_8325",+        "t_8325_8329" ~ "[](a_8326_8322, b_8326_8323, c_8326_8324)",+        "(t_8326_8327, t_8326_8328, t_8326_8329)" ~ "([]a_8326_8322, []b_8326_8323, []c_8326_8324)",+        "t_8326_8321" ~ "[](a_8327_8320, b_8327_8321, c_8327_8322)",+        "(t_8327_8325, t_8327_8326, t_8327_8327)" ~ "([]a_8327_8320, []b_8327_8321, []c_8327_8322)",+        "([]f32, []f32, []f32, []f32, []f32, []f32, []f32)" ~ "(t_8326_8327, t_8326_8328, t_8326_8329, t_8326_8320, t_8327_8325, t_8327_8326, t_8327_8327)"+      ],+      M.empty,+      M.fromList [("float_8320", (11, TyVarPrim NoLoc [FloatType Float16, FloatType Float32, FloatType Float64])), ("t_8321", (12, TyVarFree NoLoc Lifted)), ("a_8322", (13, TyVarFree NoLoc Unlifted)), ("b_8323", (13, TyVarFree NoLoc Unlifted)), ("c_8324", (13, TyVarFree NoLoc Unlifted)), ("x_8325", (13, TyVarFree NoLoc Unlifted)), ("a_8326", (13, TyVarFree NoLoc Unlifted)), ("b_8327", (13, TyVarFree NoLoc Unlifted)), ("c_8328", (13, TyVarFree NoLoc Unlifted)), ("a_8321_8325", (13, TyVarFree NoLoc Unlifted)), ("b_8321_8326", (13, TyVarFree NoLoc Unlifted)), ("c_8321_8327", (13, TyVarFree NoLoc Unlifted)), ("t_8323_8322", (14, TyVarFree NoLoc Lifted)), ("t_8324_8323", (16, TyVarFree NoLoc Lifted)), ("a_8324_8324", (17, TyVarFree NoLoc Lifted)), ("b_8324_8325", (17, TyVarFree NoLoc Lifted)), ("a_8324_8326", (17, TyVarFree NoLoc Unlifted)), ("x_8324_8327", (17, TyVarFree NoLoc Unlifted)), ("a_8325_8322", (17, TyVarFree NoLoc Unlifted)), ("b_8325_8323", (17, TyVarFree NoLoc Unlifted)), ("c_8325_8324", (17, TyVarFree NoLoc Unlifted)), ("t_8325_8329", (18, TyVarFree NoLoc Lifted)), ("t_8326_8320", (18, TyVarFree NoLoc Lifted)), ("t_8326_8321", (18, TyVarFree NoLoc Lifted)), ("a_8326_8322", (19, TyVarFree NoLoc Unlifted)), ("b_8326_8323", (19, TyVarFree NoLoc Unlifted)), ("c_8326_8324", (19, TyVarFree NoLoc Unlifted)), ("t_8326_8327", (20, TyVarFree NoLoc Lifted)), ("t_8326_8328", (20, TyVarFree NoLoc Lifted)), ("t_8326_8329", (20, TyVarFree NoLoc Lifted)), ("a_8327_8320", (21, TyVarFree NoLoc Unlifted)), ("b_8327_8321", (21, TyVarFree NoLoc Unlifted)), ("c_8327_8322", (21, TyVarFree NoLoc Unlifted)), ("t_8327_8325", (22, TyVarFree NoLoc Lifted)), ("t_8327_8326", (22, TyVarFree NoLoc Lifted)), ("t_8327_8327", (22, TyVarFree NoLoc Lifted)), ("t_8326_8327_8327_8328", (20, TyVarFree NoLoc Lifted)), ("t_8326_8328_8327_8329", (20, TyVarFree NoLoc Lifted)), ("t_8326_8329_8328_8320", (20, TyVarFree NoLoc Lifted)), ("t_8326_8320_8328_8321", (18, TyVarFree NoLoc Lifted)), ("t_8327_8325_8328_8322", (22, TyVarFree NoLoc Lifted)), ("t_8327_8326_8328_8323", (22, TyVarFree NoLoc Lifted)), ("t_8327_8327_8328_8324", (22, TyVarFree NoLoc Lifted)), ("t_8326_8321_8328_8325", (18, TyVarFree NoLoc Lifted)), ("t_8325_8329_8328_8326", (18, TyVarFree NoLoc Lifted)), ("a_8324_8324_8328_8327", (17, TyVarFree NoLoc Lifted)), ("t_8324_8323_8328_8328", (16, TyVarFree NoLoc Lifted)), ("t_8323_8322_8328_8329", (14, TyVarFree NoLoc Lifted))]+    )+  ]
src-testing/Language/Futhark/SyntaxTests.hs view
@@ -51,7 +51,7 @@  instance IsString VName where   fromString s =-    let (s', '_' : tag) = span (/= '_') s+    let (tag, s') = bimap reverse (reverse . tail) $ span (/= '_') $ reverse s      in VName (fromString s') (read tag)  pQualVName :: Parser (QualName VName)@@ -130,70 +130,74 @@         flip sizeFromName mempty <$> pQualVName       ] -pScalarNonFun :: Parser (ScalarTypeBase Size Uniqueness)-pScalarNonFun =+pScalarNonFun :: Parser d -> Parser (ScalarTypeBase d Uniqueness)+pScalarNonFun pd =   choice     [ Prim <$> pPrimType,       pTypeVar,-      tupleRecord <$> parens (pType `sepBy` lexeme ","),+      tupleRecord <$> parens (pType pd `sepBy` lexeme ","),       Record . M.fromList <$> braces (pField `sepBy1` lexeme ",")     ]   where-    pField = (,) <$> pName <* lexeme ":" <*> pType+    pField = (,) <$> pName <* lexeme ":" <*> pType pd     pTypeVar = TypeVar <$> pUniqueness <*> pQualVName <*> many pTypeArg     pTypeArg =       choice-        [ TypeArgDim <$> pSize,+        [ TypeArgDim <$> pd,           TypeArgType . second (const NoUniqueness) <$> pTypeArgType         ]     pTypeArgType =       choice         [ Scalar . Prim <$> pPrimType,-          parens pType+          parens $ pType pd         ] -pArrayType :: Parser ResType-pArrayType =+pArrayType :: Parser d -> Parser (TypeBase d Uniqueness)+pArrayType pd =   Array     <$> pUniqueness-    <*> (Shape <$> some pSize)-    <*> (second (const NoUniqueness) <$> pScalarNonFun)+    <*> (Shape <$> some pd)+    <*> (second (const NoUniqueness) <$> pScalarNonFun pd) -pNonFunType :: Parser ResType-pNonFunType =+pNonFunType :: Parser d -> Parser (TypeBase d Uniqueness)+pNonFunType pd =   choice-    [ try pArrayType,-      try $ parens pType,-      Scalar <$> pScalarNonFun+    [ try $ pArrayType pd,+      try $ parens $ pType pd,+      Scalar <$> pScalarNonFun pd     ] -pScalarType :: Parser (ScalarTypeBase Size Uniqueness)-pScalarType = choice [try pFun, pScalarNonFun]+uniquenessToDiet :: Uniqueness -> Diet+uniquenessToDiet Unique = Consume+uniquenessToDiet Nonunique = Observe++pScalarType :: Parser d -> Parser (ScalarTypeBase d Uniqueness)+pScalarType pd = choice [try pFun, pScalarNonFun pd]   where     pFun =-      pParam <* lexeme "->" <*> pRetType+      pParam <* lexeme "->" <*> pRetType pd     pParam =       choice         [ try pNamedParam,           do-            t <- pNonFunType-            pure $ Arrow Nonunique Unnamed (diet $ resToParam t) (toStruct t)+            t <- pNonFunType pd+            pure $ Arrow Nonunique Unnamed (diet $ second uniquenessToDiet t) (toStruct t)         ]     pNamedParam = parens $ do       v <- pVName <* lexeme ":"-      t <- pType-      pure $ Arrow Nonunique (Named v) (diet $ resToParam t) (toStruct t)+      t <- pType pd+      pure $ Arrow Nonunique (Named v) (diet $ second uniquenessToDiet t) (toStruct t) -pRetType :: Parser ResRetType-pRetType =+pRetType :: Parser d -> Parser (RetTypeBase d Uniqueness)+pRetType pd =   choice-    [ lexeme "?" *> (RetType <$> some (brackets pVName) <* lexeme "." <*> pType),-      RetType [] <$> pType+    [ lexeme "?" *> (RetType <$> some (brackets pVName) <* lexeme "." <*> pType pd),+      RetType [] <$> pType pd     ] -pType :: Parser ResType-pType =-  choice [try $ Scalar <$> pScalarType, pArrayType, parens pType]+pType :: Parser d -> Parser (TypeBase d Uniqueness)+pType pd =+  choice [try $ Scalar <$> pScalarType pd, pArrayType pd, parens (pType pd)]  fromStringParse :: Parser a -> String -> String -> a fromStringParse p what s =@@ -204,26 +208,40 @@  instance IsString (ScalarTypeBase Size NoUniqueness) where   fromString =-    fromStringParse (second (const NoUniqueness) <$> pScalarType) "ScalarType"+    fromStringParse+      (second (const NoUniqueness) <$> pScalarType pSize)+      "ScalarType" +instance IsString (ScalarTypeBase () NoUniqueness) where+  fromString =+    fromStringParse+      (second (const NoUniqueness) <$> pScalarType (pure ()))+      "ScalarType"++instance IsString (TypeBase () NoUniqueness) where+  fromString =+    fromStringParse+      (second (const NoUniqueness) <$> pType (brackets $ pure ()))+      "Type"+ instance IsString StructType where   fromString =-    fromStringParse (second (const NoUniqueness) <$> pType) "StructType"+    fromStringParse (second (const NoUniqueness) <$> pType pSize) "StructType"  instance IsString ParamType where   fromString =-    fromStringParse (resToParam <$> pType) "ParamType"+    fromStringParse (resToParam <$> pType pSize) "ParamType"  instance IsString ResType where   fromString =-    fromStringParse pType "ResType"+    fromStringParse (pType pSize) "ResType"  instance IsString StructRetType where   fromString =-    fromStringParse (second (pure NoUniqueness) <$> pRetType) "StructRetType"+    fromStringParse (second (pure NoUniqueness) <$> pRetType pSize) "StructRetType"  instance IsString ResRetType where-  fromString = fromStringParse pRetType "ResRetType"+  fromString = fromStringParse (pRetType pSize) "ResRetType"  instance IsString UncheckedExp where   fromString =
+ src-testing/Language/Futhark/TypeChecker/TySolveBenchmarks.hs view
@@ -0,0 +1,87 @@+module Language.Futhark.TypeChecker.TySolveBenchmarks (benchmarks) where++import Criterion (Benchmark, bench, bgroup, whnf)+import Data.Map qualified as M+import Generated.AllFutBenchmarks+import Language.Futhark (qualName)+import Language.Futhark.Syntax+import Language.Futhark.SyntaxTests ()+import Language.Futhark.TypeChecker.Constraints+  ( CtTy (..),+    Level,+    Reason (..),+    TyParams,+    TyVarInfo (..),+    TyVars,+  )+import Language.Futhark.TypeChecker.Monad (TypeError (..))+import Language.Futhark.TypeChecker.TySolve as N (Solution, UnconTyVar, solve)++(~) :: TypeBase () NoUniqueness -> TypeBase () NoUniqueness -> CtTy ()+t1 ~ t2 = CtEq (Reason mempty) t1 t2++tv :: VName -> Level -> (VName, (Level, TyVarInfo ()))+tv v lvl = (v, (lvl, TyVarFree mempty Unlifted))++solveNew ::+  ( [CtTy ()],+    TyParams,+    TyVars ()+  ) ->+  Either TypeError ([UnconTyVar], Solution)+solveNew (constraints, typarams, tyvars) = N.solve constraints typarams tyvars++generateContraints :: Int -> ([CtTy ()], TyParams, TyVars ())+generateContraints num_vars+  | num_vars <= 0 =+      ([], mempty, mempty)+  | num_vars == 1 =+      let v0_name = VName (nameFromString "v_0") 0+          ty_vars = M.fromList [tv v0_name 0]+       in ([], mempty, ty_vars)+  | otherwise =+      let var_names =+            [ VName (nameFromString ("v_" ++ show i)) i+            | i <- [0 .. num_vars - 1]+            ]++          ty_vars = M.fromList $ map (`tv` 0) var_names++          mkTy :: VName -> TypeBase () NoUniqueness+          mkTy v = Scalar (TypeVar NoUniqueness (qualName v) [])++          cts =+            zipWith+              (\v_i v_j -> mkTy v_i ~ mkTy v_j)+              (init var_names)+              (tail var_names)+              ++ ["v_0" ~ "i32"]++          ty_params = mempty+       in (cts, ty_params, ty_vars)++benchmarks :: Benchmark+benchmarks =+  bgroup+    "TySolve"+    [ bgroup "Synthetic" $+        concatMap+          ( \n ->+              [ bench (show n ++ " variables") $+                  whnf solveNew (generateContraints n)+              ]+          )+          sizes,+      bgroup "Converted" $+        concatMap+          ( \(name, dataCase) ->+              [ bench (name <> " (new)") $ whnf solveNew dataCase+              ]+          )+          allFutBenchmarkCases+    ]+  where+    start = 100+    end = 1000+    i = 100+    sizes = [start, start + i .. end]
+ src-testing/Language/Futhark/TypeChecker/TySolveTests.hs view
@@ -0,0 +1,401 @@+module Language.Futhark.TypeChecker.TySolveTests (tests) where++import Data.Loc (Loc (NoLoc))+import Data.Map qualified as M+import Futhark.Util.Pretty (docString)+import Language.Futhark.Syntax+import Language.Futhark.SyntaxTests ()+import Language.Futhark.TypeChecker.Constraints+  ( CtTy (..),+    Level,+    Reason (..),+    TyParams,+    TyVarInfo (..),+    TyVars,+  )+import Language.Futhark.TypeChecker.Monad (TypeError (TypeError), prettyTypeError)+import Language.Futhark.TypeChecker.TySolve+import Test.Tasty (TestTree, testGroup)+import Test.Tasty.HUnit (Assertion, assertBool, assertFailure, testCase, (@?=))+import Text.Regex.TDFA ((=~))++testSolve ::+  [CtTy ()] ->+  TyParams ->+  TyVars () ->+  ([UnconTyVar], Solution) ->+  Assertion+testSolve constraints typarams tyvars expected =+  case solve constraints typarams tyvars of+    Right s -> s @?= expected+    Left e -> assertFailure $ docString $ prettyTypeError e++testSolveFail ::+  [CtTy ()] ->+  TyParams ->+  TyVars () ->+  String ->+  Assertion+testSolveFail constraints typarams tyvars expected =+  case solve constraints typarams tyvars of+    Left (TypeError _ _ actualMsg) ->+      let regexMatch :: Bool = docString actualMsg =~ expected+       in assertBool "Regex doesn't match" regexMatch+    Right _ -> assertFailure "Expected type error, but got a solution"++-- When writing type variables/names here (a_0, b_1), make *sure* that+-- the numbers are distinct. These are all that actually matter for+-- determining identity.++(~) :: TypeBase () NoUniqueness -> TypeBase () NoUniqueness -> CtTy ()+t1 ~ t2 = CtEq (Reason mempty) t1 t2++tvFree :: VName -> Level -> (VName, (Level, TyVarInfo ()))+tvFree v lvl = (v, (lvl, TyVarFree mempty Unlifted))++tvRecord :: VName -> Level -> M.Map Name (TypeBase () NoUniqueness) -> (VName, (Level, TyVarInfo ()))+tvRecord v lvl fields = (v, (lvl, TyVarRecord mempty fields))++typaram :: VName -> Level -> Liftedness -> (VName, (Level, Liftedness, Loc))+typaram v lvl liftedness = (v, (lvl, liftedness, noLoc))++tests :: TestTree+tests =+  testGroup+    "Unsized constraint solver"+    [ testCase "infer unlifted" $+        testSolve+          [ "t\8320_9896" ~ "if_t\8322_9898",+            "t\8321_9897" ~ "if_t\8322_9898",+            "t\8323_9899" ~ "if_t\8322_9898"+          ]+          mempty+          ( M.fromList+              [ ("t\8320_9896", (2, TyVarFree NoLoc Lifted)),+                ("t\8321_9897", (3, TyVarFree NoLoc Lifted)),+                ("if_t\8322_9898", (4, TyVarFree NoLoc SizeLifted)),+                ("t\8323_9899", (5, TyVarFree NoLoc Lifted))+              ]+          )+          ( [("if_t\8322_9898", SizeLifted)],+            M.fromList+              [ ("t\8320_9896", Right "if_t\8322_9898"),+                ("t\8321_9897", Right "if_t\8322_9898"),+                ("t\8323_9899", Right "if_t\8322_9898")+              ]+          ),+      testCase "empty" $+        testSolve [] mempty mempty ([], mempty),+      testCase "b_1 ~ a_0" $+        testSolve+          ["b_1" ~ "a_0"]+          mempty+          (M.fromList [tvFree "b_1" 0])+          ([], M.fromList [("b_1", Right "a_0")]),+      testCase "a_0 ~ b_1" $+        testSolve+          ["a_0" ~ "b_1"]+          mempty+          (M.fromList [tvFree "a_0" 0, tvFree "b_1" 0])+          ([("b_1", Unlifted)], M.fromList [("a_0", Right "b_1")]),+      testCase "multiple" $+        testSolve+          ["b_1" ~ "a_0", "d_3" ~ "c_2", "e_4" ~ "c_2", "c_2" ~ "a_0"]+          mempty+          (M.fromList [tvFree "a_0" 0, tvFree "b_1" 0, tvFree "c_2" 0, tvFree "d_3" 0, tvFree "e_4" 0])+          ([("a_0", Unlifted)], M.fromList [("b_1", Right "a_0"), ("c_2", Right "a_0"), ("d_3", Right "a_0"), ("e_4", Right "a_0")]),+      testCase "Two variables" $+        testSolve+          ["a_0" ~ "b_1", "c_2" ~ "d_3"]+          mempty+          (M.fromList [tvFree "a_0" 0, tvFree "c_2" 0])+          ([], M.fromList [("a_0", Right "b_1"), ("c_2", Right "d_3")]),+      testCase "i32 + (i32 + i32)" $+        testSolve+          [ "i32 -> i32 -> a_0" ~ "i32 -> i32 -> i32",+            "i32 -> a_0 -> b_1" ~ "i32 -> i32 -> i32"+          ]+          mempty+          (M.fromList [tvFree "a_0" 0, tvFree "b_1" 0])+          ([], M.fromList [("a_0", Right "i32"), ("b_1", Right "i32")]),+      testCase "((λx -> λy -> x * y) i32) i32" $+        testSolve+          [ "a_0 -> b_1 -> c_2" ~ "i32 -> i32 -> i32",+            "a_0 -> b_1 -> c_2" ~ "i32 -> d_3",+            "d_3" ~ "i32 -> e_4"+          ]+          mempty+          (M.fromList [tvFree "a_0" 0, tvFree "b_1" 0, tvFree "c_2" 0, tvFree "d_3" 0, tvFree "e_4" 0])+          ( [],+            M.fromList+              [ ("a_0", Right "i32"),+                ("b_1", Right "i32"),+                ("c_2", Right "i32"),+                ("d_3", Right "i32 -> i32"),+                ("e_4", Right "i32")+              ]+          ),+      testCase "rec λf -> λn -> if n == 0 then 1 else n * (f (n - 1))" $+        testSolve+          [ "b_1 -> i32 -> c_2" ~ "i32 -> i32 -> bool",+            "b_1 -> i32 -> d_3" ~ "i32 -> i32 -> i32",+            "a_0" ~ "d_3 -> e_4",+            "b_1 -> e_4 -> f_5" ~ "i32 -> i32 -> i32",+            "c_2" ~ "bool",+            "i32" ~ "f_5",+            "g_6 -> g_6" ~ "a_0 -> b_1 -> i32"+          ]+          mempty+          (M.fromList [tvFree "a_0" 0, tvFree "b_1" 0, tvFree "c_2" 0, tvFree "d_3" 0, tvFree "e_4" 0, tvFree "f_5" 0, tvFree "g_6" 0])+          ( [],+            M.fromList+              [ ("a_0", Right "i32 -> i32"),+                ("b_1", Right "i32"),+                ("c_2", Right "bool"),+                ("d_3", Right "i32"),+                ("e_4", Right "i32"),+                ("f_5", Right "i32"),+                ("g_6", Right "i32 -> i32")+              ]+          ),+      testCase "let id = λx -> x in id id" $+        testSolve+          ["b_1 -> b_1" ~ "(c_2 -> c_2) -> d_3"]+          mempty+          (M.fromList [tvFree "b_1" 0, tvFree "c_2" 0, tvFree "d_3" 0])+          ( [("c_2", Unlifted)],+            M.fromList+              [ ("b_1", Right "c_2 -> c_2"),+                ("d_3", Right "c_2 -> c_2")+              ]+          ),+      testCase "a_0 ~ i32" $+        testSolve+          ["a_0" ~ "i32"]+          mempty+          (M.fromList [tvFree "a_0" 0])+          ([], M.fromList [("a_0", Right "i32")]),+      testCase "a_0 ~ a_0" $+        testSolve+          ["a_0" ~ "a_0"]+          mempty+          (M.fromList [tvFree "a_0" 0])+          ([("a_0", Unlifted)], mempty),+      testCase "non-unifiable types" $+        testSolveFail+          ["a_0" ~ "i32", "a_0" ~ "bool"]+          mempty+          (M.fromList [tvFree "a_0" 0])+          ".?([Cc]annot unify).?",+      testCase "infinite type (function) 1" $+        testSolveFail+          ["a_0" ~ "a_0 -> b_1"]+          mempty+          (M.fromList [tvFree "a_0" 0])+          ".?([Oo]ccurs check).?",+      -- ! This case acts weird for the original implementation.+      testCase "infinite type (function) 2" $+        testSolveFail+          ["a_0" ~ "b_1 -> i32", "b_1" ~ "c_2", "b_1" ~ "d_3", "a_0" ~ "d_3"]+          mempty+          (M.fromList [tvFree "a_0" 0, tvFree "b_1" 0, tvFree "c_2" 0, tvFree "d_3" 0])+          ".?([Oo]ccurs check).?",+      testCase "infinite type (list)" $+        testSolveFail+          ["a_0" ~ "[]a_0"]+          mempty+          (M.fromList [tvFree "a_0" 0])+          ".?([Oo]ccurs check).?",+      testCase "infinite type (tuple)" $+        testSolveFail+          ["a_0" ~ "(a_0, bool)"]+          mempty+          (M.fromList [tvFree "a_0" 0])+          ".?([Oo]ccurs check).?",+      testCase "infinite type (record) 1" $+        testSolveFail+          ["a_0" ~ "{foo: a_0, bar: f32}"]+          mempty+          (M.fromList [tvFree "a_0" 0])+          ".?([Oo]ccurs check).?",+      testCase "infinite type (record) 2" $+        testSolveFail+          ["a_0" ~ "{foo: b_1}", "b_1" ~ "c_2", "a_0" ~ "c_2"]+          mempty+          (M.fromList [tvFree "a_0" 0, tvFree "b_1" 0, tvFree "c_2" 0])+          ".?([Oo]ccurs check).?",+      testCase "infinite type (record) 3" $+        testSolveFail+          ["a_0" ~ "{foo: b_1}", "c_2" ~ "b_1", "a_0" ~ "c_2"]+          mempty+          (M.fromList [tvFree "a_0" 0, tvFree "b_1" 0, tvFree "c_2" 0])+          ".?([Oo]ccurs check).?",+      testCase "infinite type (record) 4" $+        testSolveFail+          ["a_0" ~ "{foo: b_1}", "c_2" ~ "b_1", "d_3" ~ "c_2", "a_0" ~ "c_2"]+          mempty+          (M.fromList [tvFree "a_0" 0, tvFree "b_1" 0, tvFree "c_2" 0, tvFree "d_3" 0])+          ".?([Oo]ccurs check).?",+      testCase "infinite type (consuming array param)" $+        testSolveFail+          ["a_0" ~ "*[]a_0"]+          mempty+          (M.fromList [tvFree "a_0" 0])+          ".?([Oo]ccurs check).?",+      -- ! This case acts weird for the original implementation.+      testCase "infinite type (nested)" $+        testSolveFail+          ["a_0" ~ "{foo: i32, bar: b_1}", "b_1" ~ "c_2", "c_2" ~ "i32 -> []a_0"]+          mempty+          (M.fromList [tvFree "a_0" 0, tvFree "b_1" 0, tvFree "c_2" 0])+          ".?([Oo]ccurs check).?",+      testCase "vector and 2D matrix" $+        testSolveFail+          ["a_0" ~ "[]i32", "a_0" ~ "[][]i32"]+          mempty+          (M.fromList [tvFree "a_0" 0])+          ".?([Cc]annot unify).?",+      testCase "different array types" $+        testSolveFail+          ["a_0" ~ "[]f64", "a_0" ~ "[]i64"]+          mempty+          (M.fromList [tvFree "a_0" 0])+          ".?([Cc]annot unify).?",+      testCase "simple record" $+        testSolve+          ["a_0" ~ "{foo: i32, bar: bool}"]+          mempty+          (M.fromList [tvFree "a_0" 0])+          ([], M.fromList [("a_0", Right "{foo: i32, bar: bool}")]),+      testCase "record 2" $+        testSolve+          ["a_0" ~ "{foo: b_1, bar: c_2}", "b_1" ~ "c_2", "c_2" ~ "i64"]+          mempty+          (M.fromList [tvFree "a_0" 0, tvFree "b_1" 0, tvFree "c_2" 0])+          ( [],+            M.fromList+              [ ("a_0", Right "{foo: i64, bar: i64}"),+                ("b_1", Right "i64"),+                ("c_2", Right "i64")+              ]+          ),+      testCase "record 3" $+        testSolve+          ["a_0" ~ "{foo: b_1, bar: c_2}", "b_1" ~ "c_2"]+          (M.fromList [typaram "c_2" 0 Lifted])+          (M.fromList [tvFree "a_0" 0, tvFree "b_1" 0])+          ( [],+            M.fromList+              [ ("a_0", Right "{foo: c_2, bar: c_2}"),+                ("b_1", Right "c_2")+              ]+          ),+      testCase "tuple" $+        testSolve+          ["a_0" ~ "(b_1, c_2, d_3)", "c_2" ~ "d_3"]+          mempty+          (M.fromList [tvFree "a_0" 0, tvFree "b_1" 0, tvFree "c_2" 0, tvFree "d_3" 0])+          ( [("b_1", Unlifted), ("d_3", Unlifted)],+            M.fromList+              [ ("a_0", Right "(b_1, d_3, d_3)"),+                ("c_2", Right "d_3")+              ]+          ),+      testCase "compatible levels" $+        testSolve+          ["a_0" ~ "b_1"]+          (M.fromList [typaram "a_0" 0 Unlifted])+          (M.fromList [tvFree "b_1" 1])+          ([], M.fromList [("b_1", Right "a_0")]),+      testCase "scope violation 1" $+        testSolveFail+          ["a_0" ~ "b_1"]+          (M.fromList [typaram "b_1" 1 Unlifted])+          (M.fromList [tvFree "a_0" 0])+          ".?(scope violation).?",+      testCase "scope violation 2" $+        testSolveFail+          ["a_0" ~ "b_1", "b_1" ~ "c_2"]+          (M.fromList [typaram "c_2" 1 Unlifted])+          (M.fromList [tvFree "a_0" 0, tvFree "b_1" 1])+          ".?(scope violation).?",+      testCase "differently sized tuples" $+        testSolveFail+          ["a_0" ~ "(i32, c_2)", "b_1" ~ "(i32, c_2, bool)", "a_0" ~ "b_1"]+          mempty+          (M.fromList [tvFree "a_0" 0, tvFree "b_1" 0])+          ".?([Cc]annot unify).?",+      testCase "Prim type last substitution" $+        testSolve+          [ "t\8321_8321" ~ "num\8320_8320",+            "index\8322_8322" ~ "index_elem\8323_8323",+            "[]t_0" ~ "[]index_elem\8323_8323"+          ]+          (M.fromList [typaram "t_0" 0 Unlifted])+          ( M.fromList+              [ ("num\8320_8320", (2, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64, Unsigned Int8, Unsigned Int16, Unsigned Int32, Unsigned Int64, FloatType Float16, FloatType Float32, FloatType Float64])),+                ("t\8321_8321", (2, TyVarPrim NoLoc [Signed Int8, Signed Int16, Signed Int32, Signed Int64])),+                ("index\8322_8322", (2, TyVarFree NoLoc Unlifted)),+                ("index_elem\8323_8323", (2, TyVarFree NoLoc Unlifted))+              ]+          )+          ( [],+            M.fromList+              [ ("num\8320_8320", Left [Signed Int8, Signed Int16, Signed Int32, Signed Int64]),+                ("t\8321_8321", Left [Signed Int8, Signed Int16, Signed Int32, Signed Int64]),+                ("index\8322_8322", Right "t_0"),+                ("index_elem\8323_8323", Right "t_0")+              ]+          ),+      testCase "record with polymorphic fields" $+        testSolve+          [ "d_3" ~ "{foo: e_4, bar: f_5}",+            "e_4" ~ "i32",+            "f64" ~ "f_5",+            "a_0" ~ "d_3"+          ]+          mempty+          ( M.fromList+              [ tvRecord "a_0" 0 $+                  M.fromList+                    [ ("foo", Scalar (Prim (Signed Int32))),+                      ("bar", Scalar (Prim (FloatType Float64)))+                    ],+                tvFree "d_3" 0,+                tvFree "e_4" 0,+                tvFree "f_5" 0+              ]+          )+          ( [],+            M.fromList+              [ ("a_0", Right "{foo: i32, bar: f64}"),+                ("d_3", Right "{foo: i32, bar: f64}"),+                ("e_4", Right "i32"),+                ("f_5", Right "f64")+              ]+          ),+      testCase "opaque type" $+        testSolveFail+          ["a_0" ~ "i32"]+          mempty+          mempty+          ".?([Cc]annot unify).?",+      testCase "liftedness propagation (Lifted -> SizeLifted)" $+        testSolve+          ["a_0" ~ "b_1"]+          mempty+          (M.fromList [("a_0", (0, TyVarFree mempty SizeLifted)), ("b_1", (0, TyVarFree mempty Lifted))])+          ([("b_1", SizeLifted)], M.fromList [("a_0", Right "b_1")]),+      testCase "liftedness propagation (Lifted -> Unlifted)" $+        testSolve+          ["a_0" ~ "b_1"]+          mempty+          (M.fromList [("a_0", (0, TyVarFree mempty Unlifted)), ("b_1", (0, TyVarFree mempty Lifted))])+          ([("b_1", Unlifted)], M.fromList [("a_0", Right "b_1")]),+      testCase "liftedness propagation (SizeLifted -> Unlifted)" $+        testSolve+          ["a_0" ~ "b_1"]+          mempty+          (M.fromList [("a_0", (0, TyVarFree mempty Unlifted)), ("b_1", (0, TyVarFree mempty SizeLifted))])+          ([("b_1", Unlifted)], M.fromList [("a_0", Right "b_1")])+    ]
src-testing/Language/Futhark/TypeChecker/TypesTests.hs view
@@ -33,7 +33,7 @@         assertFailure $ "Expected error, got: " <> show actual_t   where     extract (_, svars, t, _) = (svars, t)-    run = snd . runTypeM env mempty (mkInitialImport "") (newNameSource 100)+    run = snd . runTypeM env mempty (mkInitialImport "") (newNameSource 10000)     -- We hack up an environment with some predefined type     -- abbreviations for testing.  This is all pretty sensitive to the     -- specific unique names, so we have to be careful!@@ -84,64 +84,64 @@     mkNeg (x, y) = evalTest x (Left y)     pos =       [ ( "[]i32",-          ([], "?[d_100].[d_100]i32")+          ([], "?[d_10000].[d_10000]i32")         ),         ( "[][]i32",-          ([], "?[d_100][d_101].[d_100][d_101]i32")+          ([], "?[d_10000][d_10001].[d_10000][d_10001]i32")         ),         ( "bool -> []i32",-          ([], "bool -> ?[d_100].[d_100]i32")+          ([], "bool -> ?[d_10000].[d_10000]i32")         ),         ( "bool -> []f32 -> []i32",-          (["d_100"], "bool -> [d_100]f32 -> ?[d_101].[d_101]i32")+          (["d_10000"], "bool -> [d_10000]f32 -> ?[d_10001].[d_10001]i32")         ),         ( "([]i32,[]i32)",-          ([], "?[d_100][d_101].([d_100]i32, [d_101]i32)")+          ([], "?[d_10000][d_10001].([d_10000]i32, [d_10001]i32)")         ),         ( "{a:[]i32,b:[]i32}",-          ([], "?[d_100][d_101].{a:[d_100]i32, b:[d_101]i32}")+          ([], "?[d_10000][d_10001].{a:[d_10000]i32, b:[d_10001]i32}")         ),         ( "?[n].[n][n]bool",-          ([], "?[n_100].[n_100][n_100]bool")+          ([], "?[n_10000].[n_10000][n_10000]bool")         ),         ( "([]i32 -> []i32) -> bool -> []i32",-          (["d_100"], "([d_100]i32 -> ?[d_101].[d_101]i32) -> bool -> ?[d_102].[d_102]i32")+          (["d_10000"], "([d_10000]i32 -> ?[d_10001].[d_10001]i32) -> bool -> ?[d_10002].[d_10002]i32")         ),         ( "((k: i64) -> [k]i32 -> [k]i32) -> []i32 -> bool",-          (["d_101"], "((k_100: i64) -> [k_100]i32 -> [k_100]i32) -> [d_101]i32 -> bool")+          (["d_10001"], "((k_10000: i64) -> [k_10000]i32 -> [k_10000]i32) -> [d_10001]i32 -> bool")         ),         ( "square [10]",           ([], "[10][10]i32")         ),         ( "square []",-          ([], "?[d_100].[d_100][d_100]i32")+          ([], "?[d_10000].[d_10000][d_10000]i32")         ),         ( "bool -> square []",-          ([], "bool -> ?[d_100].[d_100][d_100]i32")+          ([], "bool -> ?[d_10000].[d_10000][d_10000]i32")         ),         ( "(k: i64) -> square [k]",-          ([], "(k_100: i64) -> [k_100][k_100]i32")+          ([], "(k_10000: i64) -> [k_10000][k_10000]i32")         ),         ( "fun i32 bool",           ([], "i32 -> bool")         ),         ( "fun ([]i32) bool",-          ([], "?[d_100].[d_100]i32 -> bool")+          ([], "?[d_10000].[d_10000]i32 -> bool")         ),         ( "fun bool ([]i32)",-          ([], "?[d_100].bool -> [d_100]i32")+          ([], "?[d_10000].bool -> [d_10000]i32")         ),         ( "bool -> fun ([]i32) bool",-          ([], "bool -> ?[d_100].[d_100]i32 -> bool")+          ([], "bool -> ?[d_10000].[d_10000]i32 -> bool")         ),         ( "bool -> fun bool ([]i32)",-          ([], "bool -> ?[d_100].bool -> [d_100]i32")+          ([], "bool -> ?[d_10000].bool -> [d_10000]i32")         ),         ( "pair",-          ([], "?[n_100][m_101].([n_100]i64, [m_101]i64)")+          ([], "?[n_10000][m_10001].([n_10000]i64, [m_10001]i64)")         ),         ( "(pair,pair)",-          ([], "?[n_100][m_101][n_102][m_103].(([n_100]i64, [m_101]i64), ([n_102]i64, [m_103]i64))")+          ([], "?[n_10000][m_10001][n_10002][m_10003].(([n_10000]i64, [m_10001]i64), ([n_10002]i64, [m_10003]i64))")         )       ]     neg =
src-testing/Language/Futhark/TypeCheckerTests.hs view
@@ -2,6 +2,7 @@  import Language.Futhark.TypeChecker.ConsumptionTests qualified import Language.Futhark.TypeChecker.ModulesTests qualified+import Language.Futhark.TypeChecker.TySolveTests qualified import Language.Futhark.TypeChecker.TypesTests qualified import Test.Tasty @@ -10,6 +11,7 @@   testGroup     "Source type checker tests"     [ Language.Futhark.TypeChecker.TypesTests.tests,+      Language.Futhark.TypeChecker.TySolveTests.tests,       Language.Futhark.TypeChecker.ConsumptionTests.tests,       Language.Futhark.TypeChecker.ModulesTests.tests     ]
src-testing/futhark_benchmarks.hs view
@@ -2,9 +2,11 @@  import Criterion.Main import Language.Futhark.ParserBenchmarks qualified+import Language.Futhark.TypeChecker.TySolveBenchmarks qualified  main :: IO () main =   defaultMain-    [ Language.Futhark.ParserBenchmarks.benchmarks+    [ Language.Futhark.ParserBenchmarks.benchmarks,+      Language.Futhark.TypeChecker.TySolveBenchmarks.benchmarks     ]
src/Futhark.hs view
@@ -78,12 +78,10 @@ --   optimisations such as fusion, inlining, and a host of other --   cleanup. ----- * "Futhark.IR.GPU": a representation where parallelism is expressed---   with flat /segmented operations/, and a few other GPU-specific---   operations are also supported.  The pass---   "Futhark.Pass.ExtractKernels" transforms a---   'Futhark.IR.SOACS.SOACS' program to a 'Futhark.IR.GPU.GPU'---   program.+-- * "Futhark.IR.GPU": a representation where parallelism is expressed with flat+--   /segmented operations/, and a few other GPU-specific operations are also+--   supported. The pass "Futhark.Pass.Flatten" transforms a+--   'Futhark.IR.SOACS.SOACS' program to a 'Futhark.IR.GPU.GPU' program. -- -- * "Futhark.IR.MC": a representation where parallelism is expressed --   with flat /segmented operations/, and a few other multicore-specific
src/Futhark/AD/Fwd.hs view
@@ -27,14 +27,23 @@   BasicOp $ Replicate shape $ Constant $ blankPrimValue pt zeroExp t = error $ "zeroExp: " ++ show t -tanType :: (ArrayShape s, Monoid u) => TypeBase s u -> ADM (TypeBase s u)+class (ArrayShape s) => FromShape s where+  fromShape :: Shape -> s++instance FromShape Shape where+  fromShape = id++instance FromShape ExtShape where+  fromShape = fmap Free++tanType :: (FromShape s, Monoid u) => TypeBase s u -> ADM (TypeBase s u) tanType (Acc acc ispace ts u) = do   acc_tan <- tangent acc   tan_shape <- askShape   pure $ Acc acc_tan (tan_shape <> ispace) ts u tanType t = do   shape <- askShape-  pure $ arrayOf (Prim (elemType t)) (shape `prependShape` arrayShape t) u+  pure $ arrayOf (Prim (elemType t)) (fromShape shape <> arrayShape t) u   where     u = case t of       Array _ _ u' -> u'@@ -114,7 +123,7 @@ bundleNewList :: (TanBuilder a) => [a] -> ADM [a] bundleNewList = fmap (uncurry interleave . unzip) . mapM bundleNew -instance (ArrayShape s, Monoid u) => TanBuilder (PatElem (TypeBase s u)) where+instance (FromShape s, Monoid u) => TanBuilder (PatElem (TypeBase s u)) where   newTan (PatElem p t) = do     p' <- tanVName p     insertTan p p'@@ -130,7 +139,7 @@ bundleNewPat :: (TanBuilder (PatElem t)) => Pat t -> ADM (Pat t) bundleNewPat (Pat pes) = Pat <$> bundleNewList pes -instance (ArrayShape s, Monoid u) => TanBuilder (Param (TypeBase s u)) where+instance (FromShape s, Monoid u) => TanBuilder (Param (TypeBase s u)) where   newTan (Param _ p t) = do     PatElem p' t' <- newTan $ PatElem p t     pure $ Param mempty p' t'@@ -149,7 +158,7 @@   tangent :: a -> ADM a   bundleTan :: a -> ADM (a, a) -instance (ArrayShape s, Monoid u) => Tangent (TypeBase s u) where+instance (FromShape s, Monoid u) => Tangent (TypeBase s u) where   tangent = tanType   bundleTan t = do     t' <- tangent t@@ -367,6 +376,40 @@   params' <- bundleNewList params   mkLambda params' $ bodyBind =<< fwdBody body +-- Differentiating a FlatMap lambda is mostly straightforward, except that we do+-- not care about the tangent of the size result. We must also take care to+-- reconstruct the return type, since we do not have facilities similar to+-- mkLambda for ExtLambdas. Finally, in the vector case (see Note [Forward-Mode+-- vector AD]) the tangent of a nonuniform result has the vector shape+-- outermost, but a FlatMap demands that the existential size come first, so we+-- must transpose it. It is transposed back outside the FlatMap; see 'fwdSOAC'.+fwdFlatMapLambda :: ExtLambda SOACS -> ADM (ExtLambda SOACS)+fwdFlatMapLambda (Lambda params rettype body) = do+  params' <- bundleNewList params+  (body', rettype') <- buildBody . localScope (scopeOfLParams params') $ do+    (meta_res, val_res) <- fmap (splitAt 2) <$> bodyBind =<< fwdBody body+    let (val_res_primal, val_res_tan) = unterleave val_res+        val_rettype = drop 1 rettype+    val_rettype_tan <- mapM tanType val_rettype+    (val_res_tan', val_rettype_tan') <-+      mapAndUnzipM pullExtDim $ zip3 val_rettype val_res_tan val_rettype_tan+    pure+      ( take 1 meta_res <> interleave val_res_primal val_res_tan',+        Prim int64 : interleave val_rettype val_rettype_tan'+      )+  pure $ Lambda params' rettype' body'+  where+    -- Put the existential size back in the outermost position.+    pullExtDim (t, res@(SubExpRes cs se), t_tan)+      | flatMapNonuniform t = do+          tan_shape <- askShape+          v <- pushTanShape =<< asVName se+          pure+            ( SubExpRes cs $ Var v,+              rearrangeType (vecPerm tan_shape t_tan) t_tan+            )+      | otherwise = pure (res, t_tan)+ fwdWithAccLambda :: [WithAccInput SOACS] -> Lambda SOACS -> ADM (Lambda SOACS) fwdWithAccLambda inputs (Lambda params _ body) = do   let (cert_params, acc_params) = splitAt (length inputs) params@@ -397,7 +440,7 @@     zipWithM_ (trArrParamTan tan_shape) arr_params' arr_params'_tan     (acc_res, map_res) <- fmap (splitAt (num_accs * 2)) . bodyBind =<< fwdBody body     let (map_res_primal, map_res_tan) = unterleave map_res-    map_res_tan' <- mapM (trMapResTan tan_shape) map_res_tan+    map_res_tan' <- mapM trMapResTan map_res_tan     pure $ acc_res <> interleave map_res_primal map_res_tan'   where     -- Array parameters need to be treated specially as the chunk parameter@@ -417,10 +460,8 @@       insertTan (paramName p) v      -- Put the chunk size back in the outermost position.-    trMapResTan tan_shape (SubExpRes cs ~(Var v)) = do-      v_t <- lookupType v-      let perm = vecPerm tan_shape v_t-      fmap varRes . certifying cs $ letExp (baseName v) . BasicOp $ Rearrange v perm+    trMapResTan (SubExpRes cs ~(Var v)) =+      SubExpRes cs . Var <$> pushTanShape v  pushTanShape :: VName -> ADM VName pushTanShape v = do@@ -538,6 +579,19 @@     lam_res <- auxing aux $ eLambda lam $ map eSubExp args     forM (zip (patNames pat) lam_res) $ \(v, SubExpRes cs se) ->       certifying cs $ letBindNames [v] $ BasicOp $ SubExp se+fwdSOAC (Pat pes) aux (FlatMap w arrs lam) = do+  -- We do not bother to create tangents for the metadata part of the result.+  let (meta_pes, val_pes) = splitAt 4 pes+  (Pat val_pes', to_transpose) <- soacResPat 0 0 $ Pat val_pes+  let pat' = Pat $ meta_pes <> val_pes'+  arrs' <- soacInputsWithTangents arrs+  lam' <- fwdFlatMapLambda lam+  addStm $ Let pat' aux $ Op $ FlatMap w arrs' lam'+  tan_shape <- askShape+  forM_ to_transpose $ \(rpat, v) -> do+    v_t <- lookupType v+    let perm = rearrangeInverse $ vecPerm tan_shape v_t+    letBind rpat $ BasicOp $ Rearrange v perm fwdSOAC _ _ JVP {} =   error "fwdSOAC: nested JVP not allowed." fwdSOAC _ _ VJP {} =@@ -606,11 +660,12 @@       Just (op_lam, nes) -> do         -- We assume that op_lam has unit partial derivatives (i.e., is some         -- kind of addition). This is the case for all WithAccs produced by VJP.-        lams <- mapM addLambda $ lambdaReturnType op_lam-        -- Horizontally fuse the lambdas to produce a single one.-        idx_params <- replicateM (shapeRank shape) $ newParam "idx" $ Prim int64-        let (xs, ys) = bimap concat concat $ unzip $ map (splitAt 1 . lambdaParams) lams-        op_lam' <- mkLambda (idx_params <> xs <> ys) $ mconcat <$> mapM (bodyBind . lambdaBody) lams+        -- The operator takes one index per dimension of the index space of the+        -- tangent accumulator, which includes the vector dimensions.+        op_lam' <-+          accAddLambda+            (shapeRank tan_shape + shapeRank shape)+            (lambdaReturnType op_lam)         pure $ Just (op_lam', nes)     pure (tan_shape <> shape, arrs_tan, op')   pat' <- bundleNewPat pat
src/Futhark/AD/Rev.hs view
@@ -120,7 +120,7 @@       t <- lookupType pat_adj       returnSweepCode $ do         forM_ (zip [(0 :: Int64) ..] elems) $ \(i, se) -> do-          let slice = fullSlice t [DimFix (constant i)]+          slice <- vecSlice t [DimFix (constant i)]           updateSubExpAdj se <=< letExp "elem_adj" $ BasicOp $ Index pat_adj slice     --     Index arr slice -> do@@ -136,10 +136,12 @@     Reshape arr newshape -> do       (_pat_v, pat_adj) <- commonBasicOp pat aux e m       returnSweepCode $ do+        adj_shape <- askShape         arr_shape <- arrayShape <$> lookupType arr         void $           updateAdj arr <=< letExp "adj_reshape" . BasicOp $-            Reshape pat_adj (reshapeAll (newShape newshape) arr_shape)+            Reshape pat_adj . newshapeInner adj_shape $+              reshapeAll (newShape newshape) arr_shape     --     Rearrange arr perm -> do       (_pat_v, pat_adj) <- commonBasicOp pat aux e m@@ -155,16 +157,24 @@     Replicate (Shape ns) x -> do       (_pat_v, pat_adj) <- commonBasicOp pat aux e m       returnSweepCode $ do+        adj_shape <- askShape         x_t <- subExpType x         lam <- addLambda x_t         ne <- letSubExp "zero" $ zeroExp x_t         n <- letSubExp "rep_size" =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) ns-        pat_adj_flat <--          letExp (baseName pat_adj <> "_flat") . BasicOp $-            Reshape pat_adj (reshapeAll (Shape ns) (Shape $ n : arrayDims x_t))         reduce <- reduceSOAC [Reduce Commutative lam [ne]]-        updateSubExpAdj x-          =<< letExp "rep_contrib" (Op $ Screma n [pat_adj_flat] reduce)++        contrib <- letExp "rep_contrib" <=< mapNest adj_shape (MkSolo (Var pat_adj)) $+          \(MkSolo pat_adj') -> do+            pat_adj_v <- asVName pat_adj'+            -- Flatten the replicated dimensions into a single dimension that we+            -- can reduce across.+            pat_adj_flat <-+              letExp (baseName pat_adj <> "_flat") . BasicOp . Reshape pat_adj_v $+                reshapeAll (Shape ns <> arrayShape x_t) (Shape $ n : arrayDims x_t)+            pure $ Op $ Screma n [pat_adj_flat] reduce++        updateSubExpAdj x contrib     --     Concat d (arr :| arrs) _ -> do       (_pat_v, pat_adj) <- commonBasicOp pat aux e m@@ -172,12 +182,14 @@         let sliceAdj _ [] = pure []             sliceAdj start (v : vs) = do               v_t <- lookupType v-              let w = arraySize 0 v_t+              pat_adj_t <- lookupType pat_adj+              r <- shapeRank <$> askShape+              let w = arraySize d v_t                   slice = DimSlice start w (intConst Int64 1)               pat_adj_slice <-                 letExp (baseName pat_adj <> "_slice") $                   BasicOp $-                    Index pat_adj (sliceAt v_t d [slice])+                    Index pat_adj (sliceAt pat_adj_t (r + d) [slice])               start' <- letSubExp "start" $ BasicOp $ BinOp (Add Int64 OverflowUndef) start w               slices <- sliceAdj start' vs               pure $ pat_adj_slice : slices
src/Futhark/AD/Rev/Acc.hs view
@@ -10,17 +10,29 @@ -- The general case of taking adjoints of WithAcc is tricky.  We make -- some assumptions and lay down a basic design. ----- First, we assume that any WithAccs that occur in the program are--- come from one of these sources:+-- First, we assume that any WithAccs that occur in the program come from one of+-- these sources: -- -- - A previous instance of VJP, which means we can rely on the operator having --   a constant adjoint (it's addition as appropriate to the type). -- -- - A scatter, meaning there is no operator. ----- (These can actually be distinguished by the presence of an operator, although--- we do not currently bother.)+-- This means we are in fact ignoring one potential source: --+-- - reduce_by_index_stream, where the operator is whatever the programmer+--   wrote.+--+-- This is because we assume that an operator (if it exists) is addition,+-- meaning it has a constant adjoint. This is acceptable because+-- reduce_by_index_stream is not a real part of the language, but is exposed+-- solely for testing the WithAcc machinery.+--+-- These are distinguished by the presence of an operator, which matters because+-- only the scatter overwrites: 'update_acc' on an accumulator that has an+-- operator combines with the value already in the cell, so that value keeps its+-- full sensitivity. See 'isOperatorAcc'.+-- -- Second, the adjoint of an accumulator is an array of the same type -- as the underlying array.  For example, the adjoint type of the -- primal type 'acc(c, [n], {f64})' is '[n]f64'.  In principle the@@ -31,12 +43,17 @@ -- In the return sweep, when inserting the with_acc, we still compute the -- "original" accumulator result, but modified such that its initial value is -- the adjoint of the result of the accumulator. We also modify the update_accs--- of these accumulators to be with zero values. This means that the array that--- is produced will be equal to the adjoint of the result, except for those--- places that have been updated, where it will be zero. This is intuitively--- sensible - values that have been overwritten (and so do not contribute to the--- result) should obviously have zero sensitivity.+-- of these accumulators to be with zero values. For a scatter this means that+-- the array that is produced will be equal to the adjoint of the result, except+-- for those places that have been updated, where it will be zero. This is+-- intuitively sensible - values that have been overwritten (and so do not+-- contribute to the result) should obviously have zero sensitivity. --+-- When the accumulator has an operator, nothing is overwritten, so the array+-- must instead come out as the adjoint of the result in its entirety. We get+-- that by giving the accumulator of the return sweep addition as its operator,+-- which makes those writes of zeroes leave it alone.+-- -- # Adjoint of UpdateAcc -- -- Consider primal code@@ -55,6 +72,9 @@ -- --     v_adj += acc_adj[i] --+-- and the adjoint of the accumulator going in is the adjoint of the one coming+-- out, except for a scatter, where the cell that is overwritten is zeroed.+-- -- Further, we modify the primal code so that it becomes -- --     update_acc(acc, i, 0)@@ -136,6 +156,14 @@ -- our current translation rules, they will be dead code.  As long as -- we are careful to run dead code elimination after revVJP, we should -- be good.+--+-- There is however one place where we must copy: the adjoint of the WithAcc+-- result serves double duty in the return sweep.  It is the initial value of+-- the accumulator (and thus consumed by the WithAcc), but it is *also* the+-- adjoint of the accumulator-typed result of the lambda body, and so may be+-- read inside the lambda - by the Map rule above, for example, which+-- replicates it.  These two uses are incompatible unless the WithAcc consumes+-- a copy.  -- Note [Array Adjoints of Match] --@@ -180,7 +208,9 @@ import Control.Monad import Control.Monad.Identity import Data.List ((\\))+import Data.Maybe (isJust) import Futhark.AD.Rev.Monad+import Futhark.AD.Shared (accAddLambda) import Futhark.Builder import Futhark.IR.SOACS import Futhark.Tools@@ -299,7 +329,11 @@     free_vars <- filterM isActive $ namesToList $ freeIn lam'_vec     free_accs <- filterM (fmap isAcc . lookupType) free_vars     let free_vars' = free_vars \\ free_accs-    lam'' <- diffLambda' adjs' free_vars' lam'_vec+        op_certs =+          map (paramName . fst)+            . filter (hasOperator . snd)+            $ zip (take n_inputs (lambdaParams lam'_vec)) inputs+    lam'' <- withOperatorAccs op_certs $ diffLambda' adjs' free_vars' lam'_vec     (inputs_zeroes, inputs') <-       unzip <$> zipWithM (renameInputLambda adj_shape) (chunks lengths adjs) inputs     let certs = map paramName $ take n_inputs $ lambdaParams lam''@@ -317,6 +351,7 @@     n_inputs = length inputs     lengths = map (\(_, as, _) -> length as) inputs     arrs = concatMap (\(_, as, _) -> as) inputs+    hasOperator (_, _, op) = isJust op      -- Transpose the accumulator-related adjoints from [vec...][shape...]elem     -- to [shape...][vec...]elem. Non-accumulator adjs are left unchanged.@@ -347,23 +382,37 @@               perm = [s .. s + r - 1] ++ [0 .. s - 1] ++ [s + r .. total - 1]           letExp (baseName v <> "_tr") $ BasicOp $ Rearrange v perm -    renameInputLambda adj_sh as_adj (shape, as, _) = do+    renameInputLambda adj_sh as_adj (shape, as, op) = do       -- Compute element types with vectorised dimensions included.       orig_nes_ts <- mapM (fmap (stripArray (shapeRank shape)) . lookupType) as       let vec_nes_ts = map (`arrayOfShape` adj_sh) orig_nes_ts       zeroes <- mapM (zeroArray mempty) vec_nes_ts+      -- The result adjoint is consumed by the WithAcc, but is also the adjoint+      -- of the accumulator inside the lambda (see Note [Adjoints of+      -- accumulators]), so it must remain readable there. Hence the copy.+      as' <- mapM (letExp "acc_adj_init" . BasicOp . Replicate mempty . Var <=< adjVal) as_adj       -- Transpose adjoints from [vec...][shape...]elem to [shape...][vec...]elem       -- so they match the accumulator layout.-      as' <- mapM adjVal as_adj       as'' <- mapM vecToInner as'-      pure (map Var zeroes, (shape, as'', Nothing))+      -- 'zeroOutUpdates' makes the primal updates write zeroes. For a+      -- scatter-like accumulator that overwrites the cell, this is what we+      -- want. An accumulator with a combining operator does not overwrite, so+      -- we give this one addition as its operator, making those writes no-ops.+      op' <- case op of+        Nothing -> pure Nothing+        -- Under vectorisation the element types gain the vector dimensions, so+        -- the original operator no longer fits. We assume it is addition+        -- anyway, so just build that.+        Just _ -> do+          add_lam <- accAddLambda (shapeRank shape) vec_nes_ts+          nes <- mapM (letSubExp "acc_adj_zero" . zeroExp) vec_nes_ts+          pure $ Just (add_lam, nes)+      pure (map Var zeroes, (shape, as'', op')) -    diffLambda' res_adjs get_adjs_for (Lambda params ts body) = do-      localScope (scopeOfLParams params) $ do-        Body () stms res <- vjpBody ops res_adjs get_adjs_for body-        let body' = Body () stms $ take n_inputs res <> takeLast (length get_adjs_for) res-        ts' <- mapM lookupType get_adjs_for-        pure $ Lambda params (take n_inputs ts <> ts') body'+    diffLambda' res_adjs get_adjs_for (Lambda params _ body) =+      mkLambda params $ do+        res <- bodyBind =<< vjpBody ops res_adjs get_adjs_for body+        pure $ take n_inputs res <> takeLast (length get_adjs_for) res  diffUpdateAcc ::   Pat Type ->@@ -375,22 +424,56 @@   ADM () ->   ADM () diffUpdateAcc pat aux safety acc is vs m = do+  -- By the type rules for UpdateAcc, the pattern must be a singleton.+  let Pat ~[pe] = pat   addStm $ Let pat aux $ BasicOp $ UpdateAcc safety acc is vs   m-  pat_adjs <- mapM lookupAdjVal (patNames pat)+  adj <- lookupAdjVal $ patElemName pe   returnSweepCode $ do-    forM_ (zip pat_adjs vs) $ \(adj, v) -> do-      adj_t <- lookupType adj-      let index_adj = pure $ BasicOp $ Index adj $ fullSlice adj_t $ map DimFix is-      adj_i <--        letExp "updateacc_val_adj" =<< case safety of-          Unsafe ->-            index_adj-          Safe ->-            -- The primal UpdateAcc may be out-of-bounds, in which case-            -- indexing the adjoint is dangerous.-            eIf+    adj_t <- lookupType adj+    acc_t <- lookupType acc+    -- An accumulator with a combining operator does not overwrite, so the+    -- incoming value of the updated cell retains its full sensitivity.+    overwrites <- case acc_t of+      Acc cert _ _ _ -> not <$> isOperatorAcc cert+      _ -> pure True+    let elem_t = stripArray (length is) adj_t+        slice = fullSlice adj_t $ map DimFix is+        -- The value adjoint is the corresponding cell of the accumulator+        -- adjoint.+        index_adj = maybe_copy $ pure $ BasicOp $ Index adj slice+          where+            -- We have to copy a slice because we are updating 'adj' as well -+            -- even though in many cases that update is likely dead code... Not+            -- great.+            maybe_copy+              | null $ sliceDims slice = id+              | otherwise = eCopy+        -- For a scatter-like accumulator, the input accumulator adjoint is the+        -- result adjoint with the updated cell zeroed out: a cell that is+        -- subsequently overwritten does not contribute to the result, and so+        -- has zero sensitivity.+        zeroed+          | overwrites = do+              z <- letSubExp "acc_adj_zero" $ zeroExp elem_t+              pure $ BasicOp $ Update Unsafe adj slice z+          | otherwise = pure $ BasicOp $ SubExp $ Var adj+    (adj_i, acc_adj) <- case safety of+      Unsafe ->+        (,)+          <$> (letExp "updateacc_val_adj" =<< index_adj)+          <*> (letExp "acc_adj" =<< zeroed)+      Safe -> do+        -- The primal UpdateAcc may be out-of-bounds, in which case indexing the+        -- adjoint is dangerous and the input accumulator adjoint is unchanged.+        ~[adj_i, acc_adj] <-+          letTupExp "updateacc_adj"+            =<< eIf               (eShapeInBounds (arrayShape adj_t) (map eSubExp is))-              (eBody [index_adj])-              (eBody [pure $ zeroExp $ stripArray (length is) adj_t])-      updateSubExpAdj v adj_i+              (eBody [index_adj, zeroed])+              (eBody [pure $ zeroExp elem_t, pure $ BasicOp $ SubExp $ Var adj])+        pure (adj_i, acc_adj)+    -- XXX: this is only OK because we assume accumulators are currently+    -- singleton.+    updateSubExpAdj (head vs) adj_i+    insAdj acc acc_adj
src/Futhark/AD/Rev/Map.hs view
@@ -1,12 +1,17 @@ {-# 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+-- | VJP transformation for 'Map' and 'FlatMap'. This is a pretty complicated+-- case due to the possibility of free variables. The two are handled together+-- because the return sweep of a 'FlatMap' is itself a 'Map', and so requires+-- all the same machinery.+module Futhark.AD.Rev.Map (vjpMap, vjpFlatMap) where  import Control.Monad-import Data.Bifunctor (first)+import Data.Bifunctor (first, second)+import Data.Either (rights)+import Data.Maybe (catMaybes) import Futhark.AD.Rev.Monad+import Futhark.AD.Shared (accAddLambda, asVName, vecPerm) import Futhark.Analysis.PrimExp.Convert import Futhark.Builder import Futhark.IR.SOACS@@ -75,12 +80,6 @@     subAD $ mkLambda (cert_params ++ acc_params) $ m $ map paramName acc_params   letTupExp "withhacc_res" $ WithAcc inputs acc_lam -vecPerm :: Shape -> Type -> [Int]-vecPerm adj_shape t =-  [shapeRank adj_shape]-    ++ [0 .. shapeRank adj_shape - 1]-    ++ [shapeRank adj_shape + 1 .. arrayRank t - 1]- pushAdjShape :: VName -> ADM VName pushAdjShape v = do   adj_shape <- askShape@@ -101,6 +100,128 @@       let perm = rearrangeInverse $ vecPerm adj_shape v_t       letExp (baseName v <> "_tr") $ BasicOp $ Rearrange v perm +withAccInput ::+  (VName, (Shape, PrimType)) ->+  ADM (Shape, [VName], Maybe (Lambda SOACS, [SubExp]))+withAccInput (v, (shape, pt)) = do+  v_adj <- lookupAdjVal v+  add_lam <- accAddLambda (shapeRank shape) [Prim pt]+  zero <- letSubExp "zero" $ zeroExp $ Prim pt+  pure (shape, [v_adj], Just (add_lam, [zero]))++-- | Run an action in a context where the array-typed adjoints of the given+-- free variables have been turned into accumulators, so that the contributions+-- from each iteration of the SOAC are summed.  The action is passed the free+-- variables that were given accumulator adjoints, and those that were not.  The+-- 'VName' list is the input arrays of the SOAC.+accAdjoints :: [VName] -> [VName] -> ([VName] -> Names -> ADM ()) -> ADM ()+accAdjoints as free m = do+  (arr_free, acc_free, nonacc_free) <-+    partitionAdjVars <$> classifyAdjVars free+  arr_free' <- mapM withAccInput arr_free+  -- We only consider those input arrays that are also not free in+  -- the lambda.+  let as_nonfree = filter (`notElem` free) as+  (arr_adjs, acc_adjs, rest_adjs) <-+    fmap (splitAt3 (length arr_free) (length acc_free)) . withAcc arr_free' $ \accs -> do+      zipWithM_ insAdj (map fst arr_free) accs+      () <- m (acc_free ++ map fst arr_free) (namesFromList nonacc_free)+      acc_free_adj <- mapM lookupAdjVal acc_free+      arr_free_adj <- mapM (lookupAdjVal . fst) arr_free+      nonacc_free_adj <- mapM lookupAdjVal nonacc_free+      as_nonfree_adj <- mapM lookupAdjVal as_nonfree+      pure $ varsRes $ arr_free_adj <> acc_free_adj <> nonacc_free_adj <> as_nonfree_adj+  zipWithM_ insAdj acc_free acc_adjs+  zipWithM_ insAdj (map fst arr_free) arr_adjs+  let (nonacc_adjs, as_nonfree_adjs) = splitAt (length nonacc_free) rest_adjs+  zipWithM_ insAdj nonacc_free nonacc_adjs+  zipWithM_ insAdj as_nonfree as_nonfree_adjs++-- | Add the per-iteration contributions to the adjoint of a free variable.  If+-- the adjoint is an accumulator, the summation has already taken place.+freeContrib :: SubExp -> VName -> VName -> ADM ()+freeContrib w v contribs = do+  contribs_t <- lookupType contribs+  case rowType contribs_t of+    Acc {} -> void $ insAdj v contribs+    t -> do+      lam <- addLambda t+      zero <- letSubExp "zero" $ zeroExp t+      reduce <- reduceSOAC [Reduce Commutative lam [zero]]+      contrib_sum <-+        letExp (baseName v <> "_contrib_sum") . Op $+          Screma w [contribs] reduce+      void $ updateAdj v contrib_sum++-- | Turn the 'ExtLambda' of a 'FlatMap' into an ordinary 'Lambda', by dropping+-- the size result and coercing each nonuniform result to the given size, which+-- must be dynamically equal to the size the lambda computes.  See Note+-- [Adjoints of FlatMap].+flatMapPlainLambda :: SubExp -> ExtLambda SOACS -> ADM (Lambda SOACS)+flatMapPlainLambda n (Lambda params rettype body) =+  mkLambda params $ do+    res <- bodyBind body+    forM (zip (drop 1 rettype) (drop 1 res)) $ \(t, res_i@(SubExpRes cs se)) ->+      if flatMapNonuniform t+        then do+          v <- asVName se+          v_t <- lookupType v+          fmap varRes . certifying cs . letExp "flatmap_res_coerce" $+            shapeCoerce (arrayDims (v_t `setOuterSize` n)) v+        else pure res_i++-- | Construct the Map that constitutes the return sweep of a Map-like SOAC.+-- Contributions to the free variables of the lambda are handled here; the+-- contributions to the input arrays are passed to the continuation.+mapReturnSweep ::+  (FreeIn t) =>+  VjpOps ->+  StmAux () ->+  -- | Width of the SOAC and its input arrays.+  (SubExp, [VName]) ->+  -- | Lambda of the SOAC, which must already have been renamed.  Only its+  -- parameters and free variables are used.+  GLambda SOACS t ->+  -- | Additional arrays to map across, and the parameters receiving their+  -- elements.+  [(VName, LParam SOACS)] ->+  -- | Produce the adjoints of the results of the lambda to differentiate, along+  -- with that lambda - which need not be the lambda of the SOAC, but must have+  -- the same parameters.  Run inside the return sweep lambda.+  ADM ([Adj], Lambda SOACS) ->+  -- | Given the contribution to each input array.+  ([VName] -> ADM ()) ->+  ADM ()+mapReturnSweep ops aux (w, as) lam extra mkAdjs onContribs = do+  free <- filterM isActive $ namesToList $ freeIn lam+  accAdjoints as free $ \free_with_adjs free_without_adjs -> do+    free_adjs <- mapM lookupAdjVal free_with_adjs+    free_adjs_ts <- mapM lookupType free_adjs+    free_adjs_params <- mapM (newParam "free_adj_p") free_adjs_ts+    let (extra_arrs, extra_params) = unzip extra+        adjs_for = map paramName (lambdaParams lam) ++ free+    lam_rev <-+      mkLambda (lambdaParams lam ++ extra_params ++ free_adjs_params)+        . subAD+        . noAdjsFor free_without_adjs+        $ do+          zipWithM_ insAdj free_with_adjs $ map paramName free_adjs_params+          (res_adjs, lam') <- mkAdjs+          bodyBind . lambdaBody =<< vjpLambda ops res_adjs adjs_for lam'++    (param_contribs, free_contribs) <-+      fmap (splitAt (length (lambdaParams lam))) $+        auxing aux+          . letTupExp "map_adjs"+          . Op+          . Screma w (as ++ extra_arrs ++ free_adjs)+          =<< mapSOAC lam_rev++    -- Crucial that we handle the free contribs first in case 'free'+    -- and 'as' intersect.+    zipWithM_ (freeContrib w) free free_contribs+    onContribs param_contribs+ -- | 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 ()@@ -193,81 +314,114 @@     mapM (newParam "map_adj_p" . rowType <=< lookupType) pat_adj_vals    map_lam' <- renameLambda map_lam-  free <- filterM isActive $ namesToList $ freeIn map_lam'--  accAdjoints free $ \free_with_adjs free_without_adjs -> do-    free_adjs <- mapM lookupAdjVal free_with_adjs-    free_adjs_ts <- mapM lookupType free_adjs-    free_adjs_params <- mapM (newParam "free_adj_p") free_adjs_ts-    let lam_rev_params =-          lambdaParams map_lam' ++ pat_adj_params ++ free_adjs_params-        adjs_for = map paramName (lambdaParams map_lam') ++ free-    lam_rev <--      mkLambda lam_rev_params . subAD . noAdjsFor free_without_adjs $ do-        zipWithM_ insAdj free_with_adjs $ map paramName free_adjs_params-        bodyBind . lambdaBody-          =<< vjpLambda ops (map adjFromParam pat_adj_params) adjs_for map_lam'--    (param_contribs, free_contribs) <--      fmap (splitAt (length (lambdaParams map_lam'))) $-        auxing aux-          . letTupExp "map_adjs"-          . Op-          . Screma w (as ++ pat_adj_vals ++ free_adjs)-          =<< mapSOAC lam_rev+  let param_ts = map paramType (lambdaParams map_lam')+      extra = zip pat_adj_vals pat_adj_params+      mkAdjs = pure (map adjFromParam pat_adj_params, map_lam') -    -- Crucial that we handle the free contribs first in case 'free'-    -- and 'as' intersect.-    zipWithM_ freeContrib free free_contribs-    let param_ts = map paramType (lambdaParams map_lam')+  mapReturnSweep ops aux (w, as) map_lam' extra mkAdjs $ \param_contribs ->     forM_ (zip3 param_ts as param_contribs) $ \(param_t, a, param_contrib) ->       case param_t of-        Acc {} -> freeContrib a =<< popAdjShape param_contrib -- CHECKME+        Acc {} -> freeContrib w a =<< popAdjShape param_contrib -- CHECKME         _ -> updateAdj a =<< popAdjShape param_contrib-  where-    addIdxParams n lam = do-      idxs <- replicateM n $ newParam "idx" $ Prim int64-      pure $ lam {lambdaParams = idxs ++ lambdaParams lam} -    accAddLambda n t = addIdxParams n =<< addLambda t+-- | Perform VJP on a FlatMap.  The 'Adj' list is the adjoints of the results of+-- the FlatMap, including the metadata results.  See Note [Adjoints of FlatMap].+vjpFlatMap ::+  VjpOps ->+  Pat Type ->+  [Adj] ->+  StmAux () ->+  SubExp ->+  ExtLambda SOACS ->+  [VName] ->+  ADM ()+vjpFlatMap ops pat pat_adj aux w lam as = returnSweepCode $ do+  let ((_, shape_arr, _, offset_arr), _) = flatMapSplitMeta $ patNames pat+      (_, val_adjs) = flatMapSplitMeta pat_adj -    withAccInput (v, (shape, pt)) = do-      v_adj <- lookupAdjVal v-      add_lam <- accAddLambda (shapeRank shape) $ Prim pt-      zero <- letSubExp "zero" $ zeroExp $ Prim pt-      pure (shape, [v_adj], Just (add_lam, [zero]))+  -- The size of this iteration's segment, and where in the concatenated+  -- results it begins.+  size_p <- newParam "flatmap_size_p" $ Prim int64+  offset_p <- newParam "flatmap_offset_p" $ Prim int64 -    accAdjoints free m = do-      (arr_free, acc_free, nonacc_free) <--        partitionAdjVars <$> classifyAdjVars free-      arr_free' <- mapM withAccInput arr_free-      -- We only consider those input arrays that are also not free in-      -- the lambda.-      let as_nonfree = filter (`notElem` free) as-      (arr_adjs, acc_adjs, rest_adjs) <--        fmap (splitAt3 (length arr_free) (length acc_free)) . withAcc arr_free' $ \accs -> do-          zipWithM_ insAdj (map fst arr_free) accs-          () <- m (acc_free ++ map fst arr_free) (namesFromList nonacc_free)-          acc_free_adj <- mapM lookupAdjVal acc_free-          arr_free_adj <- mapM (lookupAdjVal . fst) arr_free-          nonacc_free_adj <- mapM lookupAdjVal nonacc_free-          as_nonfree_adj <- mapM lookupAdjVal as_nonfree-          pure $ varsRes $ arr_free_adj <> acc_free_adj <> nonacc_free_adj <> as_nonfree_adj-      zipWithM_ insAdj acc_free acc_adjs-      zipWithM_ insAdj (map fst arr_free) arr_adjs-      let (nonacc_adjs, as_nonfree_adjs) = splitAt (length nonacc_free) rest_adjs-      zipWithM_ insAdj nonacc_free nonacc_adjs-      zipWithM_ insAdj as_nonfree as_nonfree_adjs+  -- 'Left' for a nonuniform result, 'Right' for a uniform one; see Note+  -- [Adjoints of FlatMap].+  (res_adjs, uniform_adj_vals) <-+    fmap (second catMaybes) . mapAndUnzipM resAdj $+      zip (drop 1 (lambdaReturnType lam)) val_adjs -    freeContrib v contribs = do-      contribs_t <- lookupType contribs-      case rowType contribs_t of-        Acc {} -> void $ insAdj v contribs-        t -> do-          lam <- addLambda t-          zero <- letSubExp "zero" $ zeroExp t-          reduce <- reduceSOAC [Reduce Commutative lam [zero]]-          contrib_sum <--            letExp (baseName v <> "_contrib_sum") . Op $-              Screma w [contribs] reduce-          void $ updateAdj v contrib_sum+  lam' <- renameLambda lam++  let extra =+        [(shape_arr, size_p), (offset_arr, offset_p)]+          ++ zip uniform_adj_vals (rights res_adjs)+      mkAdjs = do+        res_adjs' <- mapM (segmentAdj size_p offset_p) res_adjs+        plain_lam <- flatMapPlainLambda (Var $ paramName size_p) lam'+        pure (res_adjs', plain_lam)++  mapReturnSweep ops aux (w, as) lam' extra mkAdjs $ \param_contribs ->+    forM_ (zip as param_contribs) $ \(a, param_contrib) ->+      updateAdj a =<< popAdjShape param_contrib+  where+    resAdj (t, adj) = do+      adj_v <- adjVal adj+      if flatMapNonuniform t+        then pure (Left adj_v, Nothing)+        else do+          adj_v' <- pushAdjShape adj_v+          adj_p <- newParam "flatmap_res_adj_p" . rowType =<< lookupType adj_v'+          pure (Right adj_p, Just adj_v')++    segmentAdj _ _ (Right adj_p) = pure $ adjFromParam adj_p+    segmentAdj size_p offset_p (Left adj_v) = do+      adj_v_t <- lookupType adj_v+      let segment =+            DimSlice (Var $ paramName offset_p) (Var $ paramName size_p) (intConst Int64 1)+      slice <- vecSlice adj_v_t [segment]+      fmap adjFromVar . letExp (baseName adj_v <> "_slice") . BasicOp $+        Index adj_v slice++-- Note [Adjoints of FlatMap]+--+-- The return sweep of a FlatMap is an ordinary Map, not a FlatMap.  Iteration+-- 'j' of the FlatMap contributed the segment of each nonuniform (concatenated)+-- result that begins at 'offset[j]' and has length 'shape[j]', where the offset+-- and shape arrays are results of the forward sweep.  This is all we need to+-- find the part of an adjoint that a given iteration is responsible for.+--+-- The inputs are necessarily regular arrays of the same length as the Map, and+-- so are handled exactly as for a Map.  Similarly, free variables of the lambda+-- receive contributions from every iteration, and so are handled with the usual+-- accumulator machinery; see Note [Adjoints of accumulators].  The metadata+-- results are integers, so their adjoints are ignored.+--+-- The adjoints of the results are provided to the lambda in one of two ways:+--+--  * A uniform result has exactly one element per iteration, so its adjoint is+--    an ordinary Map input.+--+--  * The adjoint of a nonuniform result is the concatenation of the adjoints of+--    every segment, which cannot be split into one Map input per iteration+--    because the segments have different sizes.  Instead the entire array is+--    passed as a free variable and sliced inside the lambda.+--+-- Accumulators cannot occur among the inputs or results: the lambda would have+-- to return the updated accumulator, and an accumulator can be neither+-- concatenated nor collected into an array. (This is distinct from an array+-- whose *adjoint* is an accumulator, which happens all the time.) FIXME: it is+-- possible that one way we may return adjoints as the uniform result, but this+-- should be easy to add.+--+-- # Coercing the segment size+--+-- One wrinkle remains. Inside the lambda, a nonuniform result has type '[k]t',+-- where 'k' is the size the lambda itself computes and returns as its first+-- result. The adjoint we slice out of the concatenated adjoint has size+-- 'shape[j]' instead. These are dynamically equal, but they are distinct names+-- as far as the type checker is concerned, and the adjoint of a value of type+-- '[k]t' must have type '[k]t'. We handle this by converting the ExtLambda into+-- a Lambda whose nonuniform results are coerced to size 'shape[j]' (see+-- 'flatMapPlainLambda'), and differentiate that. Everything downstream is then+-- the ordinary machinery for differentiating a lambda, and the size mismatch is+-- dealt with by the existing rule for differentiating a coercion.
src/Futhark/AD/Rev/Monad.hs view
@@ -56,6 +56,9 @@     --     locallyNonvector,     vecToInner,+    vecSlice,+    withOperatorAccs,+    isOperatorAcc,   ) where @@ -215,7 +218,11 @@  data REnv = REnv   { envAdjShape :: Shape,-    envAttrs :: Attrs+    envAttrs :: Attrs,+    -- | The certificates of those accumulators that have a combining operator,+    -- and so do not behave like overwrites.  See Note [Adjoints of+    -- accumulators].+    envOperatorAccs :: Names   }  newtype ADM a = ADM (BuilderT SOACS (ReaderT REnv (State RState)) a)@@ -252,7 +259,7 @@     second stateNameSource $       runState         ( runReaderT (fst <$> runBuilderT m mempty) $-            REnv shape attrs+            REnv shape attrs mempty         )         (RState mempty mempty mempty vn) @@ -529,11 +536,12 @@             UpdateAcc safety v_adj' slice' [Var d']       pure v_adj'     _ -> do+      vec_slice <- vecSlice v_adj_t $ unSlice slice       v_adjslice <-         if primType t           then pure v_adj-          else letExp (baseName v <> "_slice") $ BasicOp $ Index v_adj slice-      letInPlace "updated_adj" v_adj slice =<< addExp v_adjslice d+          else letExp (baseName v <> "_slice") $ BasicOp $ Index v_adj vec_slice+      letInPlace "updated_adj" v_adj vec_slice =<< addExp v_adjslice d   insAdj v v_adj'  updateAdj :: VName -> VName -> ADM ()@@ -626,6 +634,26 @@       pure $ case v_adj of         AdjZero {} -> False         _ -> True++-- | Note that these accumulator certificates belong to accumulators that have+-- a combining operator, for the duration of the action.+withOperatorAccs :: [VName] -> ADM a -> ADM a+withOperatorAccs certs = local $ \env ->+  env {envOperatorAccs = namesFromList certs <> envOperatorAccs env}++-- | Does this accumulator certificate belong to an accumulator with a+-- combining operator, rather than a scatter-like one that overwrites?  See Note+-- [Adjoints of accumulators].+isOperatorAcc :: VName -> ADM Bool+isOperatorAcc cert = ADM $ lift $ asks $ (cert `nameIn`) . envOperatorAccs++-- | @vecSlice t is@ indexes an adjoint of type @t@ with @is@, which is written+-- against the type of the corresponding primal value.  As the vector+-- dimensions are outermost in the adjoint, they must be skipped past.+vecSlice :: Type -> [DimIndex SubExp] -> ADM (Slice SubExp)+vecSlice t is = do+  adj_shape <- askShape+  pure $ sliceAt t (shapeRank adj_shape) is  -- | If we are doing vector AD, apply 'vecPerm' to the array. vecToInner :: VName -> ADM VName
src/Futhark/AD/Rev/SOAC.hs view
@@ -202,6 +202,9 @@       (updateSubExpAdj arg <=< letExp "contrib") $         BasicOp . SubExp . resSubExp $           contrib+vjpSOAC ops pat aux soac@(FlatMap w as lam) m = do+  pat_adj <- commonSOAC pat aux soac m+  vjpFlatMap ops pat pat_adj aux w lam as vjpSOAC _ _ _ soac _ =   error $ "vjpSOAC unhandled:\n" ++ prettyString soac
src/Futhark/AD/Shared.hs view
@@ -4,18 +4,21 @@     asVName,     mapNest,     mkMap,+    accAddLambda,   ) where  import Control.Monad+import Data.Bifunctor (bimap) import Data.Foldable import Futhark.Construct import Futhark.IR.SOACS+import Futhark.Tools  -- | A permutation for transposing the vector shape past the next dimension. -- -- That is, converts @[vec...][d][elem...]@ to @[d][vec...][elem...]@.-vecPerm :: Shape -> Type -> [Int]+vecPerm :: (ArrayShape s) => Shape -> TypeBase s u -> [Int] vecPerm vec_shape t =   [shapeRank vec_shape]     ++ [0 .. shapeRank vec_shape - 1]@@ -24,6 +27,22 @@ asVName :: (MonadBuilder m) => SubExp -> m VName asVName (Var v) = pure v asVName (Constant x) = letExp "asv" $ BasicOp $ SubExp $ Constant x++-- | An addition operator for an accumulator whose index space has the given+-- rank and whose elements have the given types.  The operators for the+-- individual element types are fused horizontally, as an accumulator operator+-- must handle all of them at once.+accAddLambda ::+  (MonadBuilder m, Rep m ~ SOACS) =>+  Int ->+  [Type] ->+  m (Lambda SOACS)+accAddLambda n ts = do+  lams <- mapM addLambda ts+  idx_params <- replicateM n $ newParam "idx" $ Prim int64+  let (xs, ys) = bimap concat concat $ unzip $ map (splitAt 1 . lambdaParams) lams+  mkLambda (idx_params <> xs <> ys) $+    mconcat <$> mapM (bodyBind . lambdaBody) lams  mapNest ::   (MonadBuilder m, Rep m ~ SOACS, Traversable f) =>
src/Futhark/Analysis/Alias.hs view
@@ -132,8 +132,8 @@ analyseLambda ::   (AliasableRep rep) =>   AliasTable ->-  Lambda rep ->-  Lambda (Aliases rep)+  GLambda rep t ->+  GLambda (Aliases rep) t analyseLambda aliases lam =   let body = analyseBody aliases $ lambdaBody lam    in lam
src/Futhark/Analysis/CallGraph.hs view
@@ -15,12 +15,12 @@ import Data.Map.Strict qualified as M import Data.Maybe (isJust) import Data.Set qualified as S-import Futhark.IR.SOACS+import Futhark.IR import Futhark.Util.Pretty -type FunctionTable = M.Map Name (FunDef SOACS)+type FunctionTable rep = M.Map Name (FunDef rep) -buildFunctionTable :: Prog SOACS -> FunctionTable+buildFunctionTable :: Prog rep -> FunctionTable rep buildFunctionTable = foldl expand M.empty . progFuns   where     expand ftab f = M.insert (funDefName f) f ftab@@ -74,8 +74,8 @@ allCalledBy :: Name -> CallGraph -> S.Set Name allCalledBy f = maybe mempty fcAllCalled . M.lookup f . cgCalledByFuns --- | @buildCallGraph prog@ build the program's call graph.-buildCallGraph :: Prog SOACS -> CallGraph+-- | Build the program's call graph.+buildCallGraph :: (TraverseOpStms rep) => Prog rep -> CallGraph buildCallGraph prog =   CallGraph fg cg   where@@ -99,7 +99,7 @@  -- | @buildCallGraph ftable fg fname@ updates @fg@ with the -- contributions of function @fname@.-buildFGfun :: FunctionTable -> FunGraph -> Name -> FunGraph+buildFGfun :: (TraverseOpStms rep) => FunctionTable rep -> FunGraph -> Name -> FunGraph buildFGfun ftable fg fname =   -- Check if function is a non-builtin that we have not already   -- processed.@@ -111,23 +111,20 @@       foldl' (buildFGfun ftable) fg' $ fcAllCalled callees     _ -> fg -buildFGStms :: Stms SOACS -> FunCalls+buildFGStms :: (TraverseOpStms rep) => Stms rep -> FunCalls buildFGStms = mconcat . map buildFGstm . stmsToList -buildFGBody :: Body SOACS -> FunCalls+buildFGBody :: (TraverseOpStms rep) => Body rep -> FunCalls buildFGBody = buildFGStms . bodyStms -buildFGstm :: Stm SOACS -> FunCalls+buildFGstm :: (TraverseOpStms rep) => Stm rep -> FunCalls buildFGstm (Let (Pat (p : _)) aux (Apply fname _ _ _)) =   FunCalls (M.singleton (patElemName p) (stmAuxAttrs aux, fname)) (S.singleton fname)-buildFGstm (Let _ _ (Op op)) = execWriter $ mapSOACM folder op+buildFGstm (Let _ _ (Op op)) = execWriter $ traverseOpStms onStms op   where-    folder =-      identitySOACMapper-        { mapOnSOACLambda = \lam -> do-            tell $ buildFGBody $ lambdaBody lam-            pure lam-        }+    onStms _ stms = do+      tell $ buildFGStms stms+      pure stms buildFGstm (Let _ _ e) = execWriter $ mapExpM folder e   where     folder =
src/Futhark/Analysis/DataDependencies.hs view
@@ -118,9 +118,9 @@ -- | Determine the variables on which the results of applying -- anonymous function @lam@ to @inputs@ depend. lambdaDependencies ::-  (ASTRep rep) =>+  (ASTRep rep, FreeIn t) =>   Dependencies ->-  Lambda rep ->+  GLambda rep t ->   [Names] ->   [Names] lambdaDependencies deps lam inputs =
src/Futhark/Analysis/Metrics.hs view
@@ -17,8 +17,10 @@  import Control.Monad import Control.Monad.Writer+import Data.Foldable (for_) import Data.List (tails) import Data.Map.Strict qualified as M+import Data.Maybe (isNothing) import Data.Text (Text) import Data.Text qualified as T import Futhark.Analysis.Metrics.Type@@ -78,13 +80,15 @@     addWhat' (ctx, k) = (what : ctx, k)  -- | Compute the metrics for a program.+--+-- Metrics inside builtin functions are ignored. progMetrics :: (OpMetrics (Op rep)) => Prog rep -> AstMetrics progMetrics prog =-  actualMetrics $-    execWriter $-      runMetricsM $ do-        mapM_ funDefMetrics $ progFuns prog-        mapM_ stmMetrics $ progConsts prog+  actualMetrics . execWriter . runMetricsM $ do+    mapM_ funDefMetrics $+      filter (isNothing . isBuiltinName . funDefName) $+        progFuns prog+    mapM_ stmMetrics $ progConsts prog  funDefMetrics :: (OpMetrics (Op rep)) => FunDef rep -> MetricsM () funDefMetrics = bodyMetrics . funDefBody@@ -113,8 +117,8 @@     forM_ (zip [0 ..] cases) $ \(i, c) ->       inside (showText (i :: Int)) $ bodyMetrics $ caseBody c     inside "default" $ bodyMetrics defbody-expMetrics Apply {} =-  seen "Apply"+expMetrics (Apply fname _ _ _) =+  inside "Apply" $ for_ (isBuiltinName fname) seen expMetrics (WithAcc _ lam) =   inside "WithAcc" $ lambdaMetrics lam expMetrics (Op op) =@@ -145,5 +149,5 @@ basicOpMetrics UserParam {} = seen "UserParam"  -- | Compute metrics for this lambda.-lambdaMetrics :: (OpMetrics (Op rep)) => Lambda rep -> MetricsM ()+lambdaMetrics :: (OpMetrics (Op rep)) => GLambda rep t -> MetricsM () lambdaMetrics = bodyMetrics . lambdaBody
src/Futhark/CLI/Benchcmp.hs view
@@ -7,13 +7,15 @@ import Data.Either qualified as E import Data.List qualified as L import Data.Map qualified as M+import Data.Ord (comparing) import Data.Text qualified as T import Data.Vector qualified as V import Futhark.Bench import Futhark.Util (showText)-import Futhark.Util.Options (mainWithOptions)+import Futhark.Util.Options (mainWithOptions, optionsError) import Statistics.Sample qualified as S import System.Console.ANSI (hSupportsANSI)+import System.Console.GetOpt (ArgDescr (ReqArg), OptDescr (Option)) import System.IO (stdout) import Text.Printf (printf) @@ -286,23 +288,123 @@   putStrLn $ printf "%s%s%s%s" (header colors) (bold colors) prog (endc colors)   mapM_ (uncurry (printSpeedUp colors)) $ M.toList bench_result +-- | Geometric mean of a list of positive doubles; returns 1.0 for an empty list.+geoMean :: [Double] -> Double+geoMean [] = 1.0+geoMean xs = exp (sum (log <$> xs) / fromIntegral (length xs))++-- | Which metric to use when sorting program groups.+data SortMetric = BySignificantCount | ByGeoMeanSignificant | ByGeoMeanAll++-- | Precomputed per-group sort scores.+data GroupScore = GroupScore+  { gsSignificantCount :: Int,+    gsGeoMeanSignificant :: Double,+    gsGeoMeanAll :: Double+  }++computeGroupScore :: M.Map T.Text SpeedUp -> GroupScore+computeGroupScore datasets =+  GroupScore+    { gsSignificantCount = length sigReg,+      gsGeoMeanSignificant = geoMean (speedup <$> sigAll),+      gsGeoMeanAll = geoMean (speedup <$> all')+    }+  where+    all' = M.elems datasets+    sigAll = filter significant all'+    sigReg = filter (\s -> significant s && speedup s < 0.99) all'++-- | Whether to surface the worst or best programs first.+-- "Worst" means the most regressions / lowest speedup ratio.+data SortOrder = WorstFirst | BestFirst++scoreFor :: SortMetric -> GroupScore -> Double+scoreFor BySignificantCount gs = fromIntegral (gsSignificantCount gs)+scoreFor ByGeoMeanSignificant gs = gsGeoMeanSignificant gs+scoreFor ByGeoMeanAll gs = gsGeoMeanAll gs++-- | Sort a list of (program, datasets) pairs by the chosen metric and order.+-- Ties on the primary key fall back to geoMeanAll ascending then name alphabetically.+sortGroups ::+  SortMetric ->+  SortOrder ->+  [(T.Text, M.Map T.Text SpeedUp)] ->+  [(T.Text, M.Map T.Text SpeedUp)]+sortGroups metric order = L.sortBy cmpFull+  where+    score g = scoreFor metric (computeGroupScore (snd g))+    geoAll g = gsGeoMeanAll (computeGroupScore (snd g))+    -- For count metrics "worst first" = descending; for ratio metrics+    -- "worst first" = ascending (lower ratio = bigger regression).+    primaryCmp = case (metric, order) of+      (BySignificantCount, WorstFirst) -> flip compare `on` score+      (BySignificantCount, BestFirst) -> compare `on` score+      (_, WorstFirst) -> compare `on` score+      (_, BestFirst) -> flip compare `on` score+    cmpFull a b =+      primaryCmp a b+        <> comparing geoAll a b+        <> comparing fst a b+    on f g x y = f (g x) (g y)++-- | Config for the benchcmp tool.+data BenchcmpConfig = BenchcmpConfig+  { cfgSortMetric :: Maybe SortMetric,+    cfgSortOrder :: SortOrder+  }++defaultConfig :: BenchcmpConfig+defaultConfig = BenchcmpConfig Nothing WorstFirst+ -- | Given a Map of programs with dataset speedups and relevant errors, print -- the errors and print the speedups in a human readable manner. printComparisons ::   Colors ->+  BenchcmpConfig ->   M.Map T.Text (M.Map T.Text SpeedUp) ->   ([T.Text], [T.Text]) ->   IO ()-printComparisons colors speedups (errors, missing) = do+printComparisons colors cfg speedups (errors, missing) = do   mapM_ (putStrLn . T.unpack) $ L.sort missing   mapM_ (putStrLn . T.unpack) $ L.sort errors-  mapM_ (uncurry (printProgSpeedUps colors)) $ M.toList speedups+  let groups = case cfgSortMetric cfg of+        Nothing -> M.toList speedups+        Just metric -> sortGroups metric (cfgSortOrder cfg) (M.toList speedups)+  mapM_ (uncurry (printProgSpeedUps colors)) groups  -- | Run @futhark benchcmp@ main :: String -> [String] -> IO ()-main = mainWithOptions () [] "<file> <file>" f+main = mainWithOptions defaultConfig options "<file> <file>" f   where-    f [a_path', b_path'] () = Just $ do+    options =+      [ Option+          []+          ["sort-by"]+          ( ReqArg+              ( \arg -> case arg of+                  "significant" -> Right $ \cfg -> cfg {cfgSortMetric = Just BySignificantCount}+                  "geomean-significant" -> Right $ \cfg -> cfg {cfgSortMetric = Just ByGeoMeanSignificant}+                  "geomean-all" -> Right $ \cfg -> cfg {cfgSortMetric = Just ByGeoMeanAll}+                  _ -> Left . optionsError $ "Unknown --sort-by value: " <> arg+              )+              "METRIC"+          )+          "Sort program groups by: significant, geomean-significant, geomean-all (default: unsorted).",+        Option+          []+          ["order"]+          ( ReqArg+              ( \arg -> case arg of+                  "worst-first" -> Right $ \cfg -> cfg {cfgSortOrder = WorstFirst}+                  "best-first" -> Right $ \cfg -> cfg {cfgSortOrder = BestFirst}+                  _ -> Left . optionsError $ "Unknown --order value: " <> arg+              )+              "ORDER"+          )+          "Sort order: worst-first (default) or best-first."+      ]+    f [a_path', b_path'] cfg = Just $ do       let a_path = T.pack a_path'       let b_path = T.pack b_path'       a_either <- decodeFileBenchResultsMap a_path@@ -316,7 +418,7 @@               else nonTtyColors        let comparePrint =-            (uncurry (printComparisons colors) .)+            (uncurry (printComparisons colors cfg) .)               . compareBenchResults a_path b_path        case (a_either, b_either) of
src/Futhark/CLI/Dev.hs view
@@ -53,11 +53,12 @@ import Futhark.Pass.ExplicitAllocations.GPU qualified as GPU import Futhark.Pass.ExplicitAllocations.MC qualified as MC import Futhark.Pass.ExplicitAllocations.Seq qualified as Seq-import Futhark.Pass.ExtractKernels import Futhark.Pass.ExtractMulticore import Futhark.Pass.FirstOrderTransform+import Futhark.Pass.Flatten (flattenSOACs) import Futhark.Pass.LiftAllocations as LiftAllocations import Futhark.Pass.LowerAllocations as LowerAllocations+import Futhark.Pass.NoGrid import Futhark.Pass.Simplify import Futhark.Passes import Futhark.Util.Log@@ -721,13 +722,14 @@     kernelsPassOption optimiseArrayLayoutGPU [],     mcPassOption optimiseArrayLayoutMC [],     kernelsPassOption addGlobalParams [],+    kernelsPassOption noGrid [],     kernelsPassOption optimiseGenRed [],     kernelsPassOption tileLoops [],     kernelsPassOption histAccsGPU [],     unstreamOption [],     sinkOption [],     kernelsPassOption reduceDeviceSyncs [],-    typedPassOption soacsProg GPU extractKernels [],+    typedPassOption soacsProg GPU flattenSOACs [],     typedPassOption soacsProg MC extractMulticore [],     allocateOption "a",     kernelsMemPassOption doubleBufferGPU [],
src/Futhark/CodeGen/Backends/GenericPython.hs view
@@ -1413,9 +1413,11 @@ compileCode (Imp.Copy t shape (dst, dstspace) (dstoffset, dststrides) (src, srcspace) (srcoffset, srcstrides)) = do   cp <- asks $ M.lookup (dstspace, srcspace) . opsCopies . envOperations   case cp of-    Nothing ->-      compileCopy t shape (dst, dstspace) (dstoffset, dststrides) (src, srcspace) (srcoffset, srcstrides)-    Just cp' -> do+    -- Values of unit type occupy no memory, so the arrays involved may well be+    -- zero-sized. The specialised copies compute the element size from the+    -- type, which they cannot do for 'Unit', so let the generic copy handle it+    -- (where it degenerates to writing nothing).+    Just cp' | t /= Unit -> do       shape' <- traverse (traverse (compileExp . untyped)) shape       dst' <- compileVar dst       src' <- compileVar src@@ -1424,6 +1426,8 @@       srcoffset' <- traverse (compileExp . untyped) srcoffset       srcstrides' <- traverse (traverse (compileExp . untyped)) srcstrides       cp' t shape' dst' (dstoffset', dststrides') src' (srcoffset', srcstrides')+    _ ->+      compileCopy t shape (dst, dstspace) (dstoffset, dststrides) (src, srcspace) (srcoffset, srcstrides) compileCode (Imp.Write dst (Imp.Count idx) pt space _ elemexp) = do   dst' <- compileVar dst   idx' <- compileExp $ Imp.untyped idx
src/Futhark/CodeGen/Backends/GenericPython/AST.hs view
@@ -122,52 +122,46 @@   pretty (Lambda p e) = "lambda" <+> pretty p <> ":" <+> pretty e   pretty None = "None" +-- | The indented body of a compound statement. Python has no empty block, so+-- an empty body must be printed as @pass@.+pyBlock :: [PyStmt] -> Doc a+pyBlock [] = indent 2 "pass"+pyBlock stms = indent 2 $ stack $ map pretty stms+ instance Pretty PyStmt where-  pretty (If cond [] []) =-    "if"-      <+> pretty cond-      <> ":"-        </> indent 2 "pass"-  pretty (If cond [] fbranch) =-    "if"-      <+> pretty cond-      <> ":"-        </> indent 2 "pass"-        </> "else:"-        </> indent 2 (stack $ map pretty fbranch)   pretty (If cond tbranch []) =     "if"       <+> pretty cond       <> ":"-        </> indent 2 (stack $ map pretty tbranch)+        </> pyBlock tbranch   pretty (If cond tbranch fbranch) =     "if"       <+> pretty cond       <> ":"-        </> indent 2 (stack $ map pretty tbranch)+        </> pyBlock tbranch         </> "else:"-        </> indent 2 (stack $ map pretty fbranch)+        </> pyBlock fbranch   pretty (Try pystms pyexcepts) =     "try:"-      </> indent 2 (stack $ map pretty pystms)+      </> pyBlock pystms       </> stack (map pretty pyexcepts)   pretty (While cond body) =     "while"       <+> pretty cond       <> ":"-        </> indent 2 (stack $ map pretty body)+        </> pyBlock body   pretty (For i what body) =     "for"       <+> pretty i       <+> "in"       <+> pretty what       <> ":"-        </> indent 2 (stack $ map pretty body)+        </> pyBlock body   pretty (With what body) =     "with"       <+> pretty what       <> ":"-        </> indent 2 (stack $ map pretty body)+        </> pyBlock body   pretty (Assign e1 e2) = pretty e1 <+> "=" <+> pretty e2   pretty (AssignOp op e1 e2) = pretty e1 <+> pretty (op ++ "=") <+> pretty e2   pretty (Comment s body) = "#" <> pretty s </> stack (map pretty body)
src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs view
@@ -97,7 +97,7 @@ -- | The maximum number of operators we support in a single SegRed. -- This limit arises out of the static allocation of counters. maxNumOps :: Int-maxNumOps = 20+maxNumOps = 25  -- | Code generation for the body of the SegRed, taking a continuation -- for saving the results of the body.  The results should be@@ -334,7 +334,7 @@       global_tid = Imp.le64 $ segFlat space       n = pe64 $ last dims -  counters <- genZeroes "counters" maxNumOps+  counters <- genZeroes "counters" $ length segbinops    reds_block_res_arrs <- groupResultArrays num_tblocks_se tblock_size_se segbinops @@ -551,8 +551,9 @@   -- anywhere?  There are other places in the compiler that will fail   -- if the block count exceeds the maximum block size, which is at   -- most 1024 anyway.-  let num_counters = maxNumOps * 1024-  counters <- genZeroes "counters" $ fromIntegral num_counters+  let counters_per_op = 1024+  counters <-+    genZeroes "counters" $ fromIntegral $ length segbinops * counters_per_op    let attrs =         (defKernelAttrs num_tblocks tblock_size)@@ -610,9 +611,9 @@             forM_ (zip4 segred_pess slugs new_lambdas [0 ..]) $               \(pes, slug, new_lambda, i) -> do                 let counter_idx =-                      fromIntegral (i * num_counters)+                      fromIntegral (i * counters_per_op)                         + flat_segment_id-                          `rem` fromIntegral num_counters+                          `rem` fromIntegral counters_per_op                 reductionStageTwo                   pes                   virttblock_id
src/Futhark/Construct.hs view
@@ -107,7 +107,7 @@     sliceDim,     fullSlice,     fullSliceNum,-    isFullSlice,+    isIdentitySlice,     sliceAt,     iota64, @@ -594,14 +594,17 @@ fullSliceNum dims slice =   Slice $ slice ++ map (\d -> DimSlice 0 d 1) (drop (length slice) dims) --- | Does the slice describe the full size of the array?  The most--- obvious such slice is one that 'DimSlice's the full span of every--- dimension, but also one that fixes all unit dimensions.-isFullSlice :: Shape -> Slice SubExp -> Bool-isFullSlice shape slice = and $ zipWith allOfIt (shapeDims shape) (unSlice slice)+-- | Does indexing an array of this shape with this slice produce the array+-- itself? This is a slice that selects every element (which may be fixing unit+-- dimensions) and has zero offset and unit stride.+isIdentitySlice :: Shape -> Slice SubExp -> Bool+isIdentitySlice shape slice =+  length (shapeDims shape) == length (unSlice slice)+    && and (zipWith allOfIt (shapeDims shape) (unSlice slice))   where     allOfIt (Constant v) DimFix {} = oneIsh v-    allOfIt d (DimSlice _ n _) = d == n+    allOfIt d (DimSlice (Constant off) n (Constant s)) =+      d == n && zeroIsh off && oneIsh s     allOfIt _ _ = False  -- | Conveniently construct a body that contains no bindings.
src/Futhark/IR/GPU.hs view
@@ -63,6 +63,9 @@   asSegOp _ = Nothing   segOp = SegOp +instance TraverseOpStms GPU where+  traverseOpStms = traverseHostOpStms traverseSOACStms+ -- Note [GPU Terminology] -- -- For lack of a better spot to put it, this Note summarises the
src/Futhark/IR/GPU/Simplify.hs view
@@ -65,7 +65,13 @@ simplifyKernelOp _ (GPUBody ts body) = do   ts' <- Engine.simplify ts   (hoisted, body') <--    Engine.simplifyBody keepOnGPU mempty (map (const mempty) ts) body+    -- A GPUBody is a single-threaded kernel that cannot perform+    -- allocations.  By pretending we are simplifying memory (as we do+    -- inside SegOps) we allow allocations to be hoisted out of+    -- branches within the body, and ultimately out of the GPUBody+    -- itself.+    Engine.localVtable (\vtable -> vtable {ST.simplifyMemory = True}) $+      Engine.simplifyBody keepOnGPU mempty (map (const mempty) ts) body   pure (GPUBody ts' body', hoisted)   where     keepOnGPU _ _ = keepExpOnGPU . stmExp
src/Futhark/IR/Parse.hs view
@@ -154,6 +154,9 @@ pExtType :: Parser ExtType pExtType = pTypeBase pExtShape (pure NoUniqueness) +pExtTypes :: Parser [ExtType]+pExtTypes = braces $ pExtType `sepBy` pComma+ pRank :: Parser Rank pRank = Rank . length <$> many (lexeme "[" *> lexeme "]") @@ -556,6 +559,19 @@       keyword "nilFn" $> Lambda mempty [] (Body (pBodyDec pr) mempty [])     ] +pExtLambda :: PR rep -> Parser (ExtLambda rep)+pExtLambda pr =+  choice+    [ lexeme "\\"+        $> Lambda+        <*> pLParams pr+        <* pColon+        <*> pExtTypes+        <* pArrow+        <*> pBody pr,+      keyword "nilFn" $> Lambda mempty [] (Body (pBodyDec pr) mempty [])+    ]+ pReduce :: PR rep -> Parser (SOAC.Reduce rep) pReduce pr =   SOAC.Reduce@@ -757,6 +773,7 @@       keyword "redomap" *> pScrema pRedomapForm,       keyword "scanomap" *> pScrema pScanomapForm,       keyword "screma" *> pScrema pScremaForm,+      keyword "flatmap" *> pFlatMap,       keyword "vjp" *> pVJP,       keyword "jvp" *> pJVP,       pHist,@@ -829,6 +846,14 @@             <*> braces (pSubExp `sepBy` pComma)             <* pComma             <*> pLambda pr+    pFlatMap =+      parens $+        SOAC.FlatMap+          <$> pSubExp+          <* pComma+          <*> braces (pVName `sepBy` pComma)+          <* pComma+          <*> pExtLambda pr     pStream = keyword "streamSeq" *> pStreamSeq     pStreamSeq =       parens $
src/Futhark/IR/Pretty.hs view
@@ -7,6 +7,7 @@   ( prettyTuple,     prettyTupleLines,     prettyString,+    prettyRet,     PrettyRep (..),   ) where@@ -376,7 +377,7 @@                   comma </> parens (pretty op' <> comma </> ppTuple' (map pretty nes))           ) -instance (PrettyRep rep) => Pretty (Lambda rep) where+instance (PrettyRep rep, Pretty t) => Pretty (GLambda rep t) where   pretty (Lambda [] [] (Body _ stms [])) | stms == mempty = "nilFn"   pretty (Lambda params rettype body) =     "\\"
src/Futhark/IR/Prop.hs view
@@ -34,6 +34,8 @@     lamIsBinOp,     isIdentityLambda,     isNilLambda,+    builtinName,+    isBuiltinName,     ASTConstraints,     IsOp (..),     ASTRep (..),@@ -46,6 +48,7 @@ import Data.Map.Strict qualified as M import Data.Maybe (isJust, mapMaybe) import Data.Set qualified as S+import Data.Text qualified as T import Futhark.IR.Pretty import Futhark.IR.Prop.Constants import Futhark.IR.Prop.Names@@ -76,10 +79,28 @@ asBasicOp (BasicOp op) = Just op asBasicOp _ = Nothing --- | An expression is safe if it is always well-defined (assuming that--- any required certificates have been checked) in any context.  For--- example, array indexing is not safe, as the index may be out of--- bounds.  On the other hand, adding two numbers cannot fail.+-- | An expression is safe if it is always well-defined in any context, assuming+-- data dependencies are satisfied - and note that certificates are just a+-- special case of data dependencies. For example, array indexing is not safe,+-- as the index may be out of bounds. Array slicing is also considered unsafe,+-- as semantically we consider this to be a copy (even if operationally it is+-- usually a view). On the other hand, adding two numbers cannot fail.+--+-- The certificate requirement is subtle but important: Iota is safe only when+-- the argument given is non-negative, which is assumed checked by a+-- certificate.+--+-- Essentially the rule is this: if an expression is safe, then it means we can+-- hoist it out of control flow without a soundness issue.+--+-- Some of these operations are less safe than you might intuitively expect,+-- e.g. Reshape is considered unsafe despite usually not resulting in any code+-- generation. Reshape can *semantically* fail if the original array does not+-- have the same number of elements as the new shape.+--+-- Being very generous with unsafety-status is to ease program transformation,+-- such that we do not have to worry too much about turning unsafe operations+-- into "safe" operations. safeExp :: (ASTRep rep) => Exp rep -> Bool safeExp (BasicOp op) = safeBasicOp op   where@@ -117,12 +138,10 @@     safeBasicOp ConvOp {} = True     safeBasicOp Scratch {} = True     safeBasicOp Concat {} = True-    safeBasicOp Reshape {} = True     safeBasicOp Rearrange {} = True     safeBasicOp Manifest {} = True     safeBasicOp Iota {} = True     safeBasicOp Replicate {} = True-    safeBasicOp (Index _ slice) = sliceShape slice /= mempty     safeBasicOp _ = False safeExp (Loop _ _ body) = safeBody body safeExp (Apply fname _ _ _) =@@ -296,3 +315,12 @@ isNilLambda lam =   null (lambdaParams lam)     && isIdentityLambda lam++-- | Construct a name for a builtin function. This name will be recognised by+-- 'isBuiltinName'.+builtinName :: T.Text -> Name+builtinName = nameFromText . ("builtin/" <>)++-- | Is this a builtin name, and if so, what is the underlying name?+isBuiltinName :: Name -> Maybe T.Text+isBuiltinName v = "builtin/" `T.stripPrefix` nameToText v
src/Futhark/IR/Prop/Names.hs view
@@ -254,9 +254,10 @@     FreeIn (LetDec rep),     FreeIn (RetType rep),     FreeIn (BranchType rep),-    FreeIn (Op rep)+    FreeIn (Op rep),+    FreeIn t   ) =>-  FreeIn (Lambda rep)+  FreeIn (GLambda rep t)   where   freeIn' (Lambda params body rettype) =     fvBind (namesFromList $ map paramName params) $@@ -447,5 +448,5 @@ boundByStms = foldMap boundByStm  -- | The names of the lambda parameters plus the index parameter.-boundByLambda :: Lambda rep -> [VName]+boundByLambda :: GLambda rep t -> [VName] boundByLambda lam = map paramName (lambdaParams lam)
src/Futhark/IR/Prop/Types.hs view
@@ -25,7 +25,6 @@     peelArray,     stripArray,     arrayDims,-    arrayExtDims,     shapeSize,     arraySize,     arraysSize,@@ -36,8 +35,6 @@     mapOnExtType,     mapOnType,     diet,-    subtypeOf,-    subtypesOf,     toDecl,     fromDecl,     isExt,@@ -188,7 +185,6 @@ -- size is the given dimension.  This is just a convenient wrapper -- around 'arrayOf'. arrayOfRow ::-  (ArrayShape (ShapeBase d)) =>   TypeBase (ShapeBase d) NoUniqueness ->   d ->   TypeBase (ShapeBase d) NoUniqueness@@ -208,7 +204,6 @@ -- | Replace the size of the outermost dimension of an array.  If the -- given type is not an array, it is returned unchanged. setOuterSize ::-  (ArrayShape (ShapeBase d)) =>   TypeBase (ShapeBase d) u ->   d ->   TypeBase (ShapeBase d) u@@ -217,7 +212,6 @@ -- | Replace the size of the given dimension of an array.  If the -- given type is not an array, it is returned unchanged. setDimSize ::-  (ArrayShape (ShapeBase d)) =>   Int ->   TypeBase (ShapeBase d) u ->   d ->@@ -249,7 +243,7 @@ -- | @stripArray n t@ removes the @n@ outermost layers of the array. -- Essentially, it is the type of indexing an array of type @t@ with -- @n@ indexes.-stripArray :: Int -> TypeBase Shape u -> TypeBase Shape u+stripArray :: Int -> TypeBase (ShapeBase d) u -> TypeBase (ShapeBase d) u stripArray n (Array et shape u)   | n < shapeRank shape = Array et (stripDims n shape) u   | otherwise = Prim et@@ -264,14 +258,9 @@  -- | Return the dimensions of a type - for non-arrays, this is the -- empty list.-arrayDims :: TypeBase Shape u -> [SubExp]+arrayDims :: TypeBase (ShapeBase d) u -> [d] arrayDims = shapeDims . arrayShape --- | Return the existential dimensions of a type - for non-arrays,--- this is the empty list.-arrayExtDims :: TypeBase ExtShape u -> [ExtSize]-arrayExtDims = shapeDims . arrayShape- -- | Return the size of the given dimension.  If the dimension does -- not exist, the zero constant is returned. arraySize :: Int -> TypeBase Shape u -> SubExp@@ -284,9 +273,9 @@ arraysSize _ [] = constant (0 :: Int64) arraysSize i (t : _) = arraySize i t --- | Return the immediate row-type of an array.  For @[[int]]@, this--- would be @[int]@.-rowType :: TypeBase Shape u -> TypeBase Shape u+-- | Return the immediate row-type of an array.  For @[][]t@, this+-- would be @[]t@.+rowType :: TypeBase (ShapeBase d) u -> TypeBase (ShapeBase d) u rowType = stripArray 1  -- | A type is a primitive type if it is not an array or memory block.@@ -319,7 +308,7 @@ -- | Rearrange the dimensions of the type.  If the length of the -- permutation does not match the rank of the type, the permutation -- will be extended with identity.-rearrangeType :: [Int] -> Type -> Type+rearrangeType :: [Int] -> TypeBase (ShapeBase d) u -> TypeBase (ShapeBase d) u rearrangeType perm t =   t `setArrayShape` Shape (rearrangeShape perm' $ arrayDims t)   where@@ -372,41 +361,13 @@ -- | @diet t@ returns a description of how a function parameter of -- type @t@ might consume its argument. diet :: TypeBase shape Uniqueness -> Diet-diet Prim {} = ObservePrim+diet Prim {} = Observe diet (Acc _ _ _ Unique) = Consume diet (Acc _ _ _ Nonunique) = Observe diet (Array _ _ Unique) = Consume diet (Array _ _ Nonunique) = Observe diet Mem {} = Observe --- | @x \`subtypeOf\` y@ is true if @x@ is a subtype of @y@ (or equal to--- @y@), meaning @x@ is valid whenever @y@ is.-subtypeOf ::-  (Ord u, ArrayShape shape) =>-  TypeBase shape u ->-  TypeBase shape u ->-  Bool-subtypeOf (Array t1 shape1 u1) (Array t2 shape2 u2) =-  u2-    <= u1-    && t1-      == t2-    && shape1-      `subShapeOf` shape2-subtypeOf t1 t2 = t1 == t2---- | @xs \`subtypesOf\` ys@ is true if @xs@ is the same size as @ys@,--- and each element in @xs@ is a subtype of the corresponding element--- in @ys@..-subtypesOf ::-  (Ord u, ArrayShape shape) =>-  [TypeBase shape u] ->-  [TypeBase shape u] ->-  Bool-subtypesOf xs ys =-  length xs == length ys-    && and (zipWith subtypeOf xs ys)- -- | Add the given uniqueness information to the types. toDecl ::   TypeBase shape NoUniqueness ->@@ -468,8 +429,8 @@ hasStaticShape (Array bt (Shape shape) u) =   Array bt <$> (Shape <$> mapM isFree shape) <*> pure u --- | Given two lists of 'ExtType's of the same length, return a list--- of 'ExtType's that is a subtype of the two operands.+-- | Given two lists of 'ExtType's of the same length, return a list of+-- 'ExtType's that generalises the two operands. generaliseExtTypes ::   [TypeBase ExtShape u] ->   [TypeBase ExtShape u] ->@@ -514,7 +475,7 @@  -- | Produce a mapping for the dimensions context. shapeExtMapping :: [TypeBase ExtShape u] -> [TypeBase Shape u1] -> M.Map Int SubExp-shapeExtMapping = dimMapping arrayExtDims arrayDims match mappend+shapeExtMapping = dimMapping arrayDims arrayDims match mappend   where     match Free {} _ = mempty     match (Ext i) dim = M.singleton i dim
src/Futhark/IR/Rephrase.hs view
@@ -89,7 +89,7 @@     <*> pure res  -- | Rephrase a lambda.-rephraseLambda :: (Monad m) => Rephraser m from to -> Lambda from -> m (Lambda to)+rephraseLambda :: (Monad m) => Rephraser m from to -> GLambda from t -> m (GLambda to t) rephraseLambda rephraser lam = do   body' <- rephraseBody rephraser $ lambdaBody lam   params' <- mapM (rephraseParam $ rephraseLParamDec rephraser) $ lambdaParams lam
src/Futhark/IR/RetType.hs view
@@ -61,11 +61,7 @@    applyRetType extret params args =     if length args == length params-      && and-        ( zipWith subtypeOf argtypes $-            expectedTypes (map paramName params) params $-              map fst args-        )+      && argtypes == expectedTypes (map paramName params) params (map fst args)       then Just $ map correctExtDims extret       else Nothing     where
src/Futhark/IR/SOACS.hs view
@@ -45,16 +45,21 @@  instance PrettyRep SOACS +instance TraverseOpStms SOACS where+  traverseOpStms = traverseSOACStms+ usesAD :: Prog SOACS -> Bool usesAD prog = any stmUsesAD (progConsts prog) || any funUsesAD (progFuns prog)   where     funUsesAD = bodyUsesAD . funDefBody     bodyUsesAD = any stmUsesAD . bodyStms     stmUsesAD = expUsesAD . stmExp+    lamUsesAD :: GLambda SOACS t -> Bool     lamUsesAD = bodyUsesAD . lambdaBody     expUsesAD (Op JVP {}) = True     expUsesAD (Op VJP {}) = True     expUsesAD (Op WithVJP {}) = True+    expUsesAD (Op (FlatMap _ _ lam)) = lamUsesAD lam     expUsesAD (Op (Stream _ _ _ lam)) = lamUsesAD lam     expUsesAD (Op (Screma _ _ (ScremaForm lam scans reds post_lam))) =       lamUsesAD lam
src/Futhark/IR/SOACS/SOAC.hs view
@@ -18,6 +18,11 @@     composeBinds,     scremaType,     soacType,+    flatMapNonuniform,+    flatMapRowTypes,+    flatMapUniformTypes,+    flatMapSplitMeta,+    flatMapSplitValues,     typeCheckSOAC,     mkIdentityLambda,     nilFn,@@ -50,6 +55,7 @@ import Control.Monad.Identity import Control.Monad.State.Strict import Control.Monad.Writer+import Data.Either (partitionEithers) import Data.List (intersperse) import Data.Map.Strict qualified as M import Data.Maybe@@ -87,6 +93,30 @@   | -- | A combination of scan, reduction, and map.  The first     -- t'SubExp' is the size of the input arrays.     Screma SubExp [VName] (ScremaForm rep)+  | -- | Nonuniform map where the nonuniform results are implicitly concatenated.+    -- The 'ExtLambda' first returns a size @k@, then its value results, each of+    -- which is either /nonuniform/ (an array with @k@ as its outermost size) or+    -- /uniform/ (not mentioning @k@ at all). The size @k@ may be used *only* as+    -- an outermost dimension.+    --+    -- For input of length @n@ returns the following:+    --+    -- * An integer @m@ denoting the size of the data arrays.+    --+    -- * The "shape array" of type @[n]i64@, giving the size of each segment.+    --   This array sums to @m@.+    --+    -- * The "flag array" of type @[m]bool@, indicating for each element when a+    --   new segment begins.+    --+    -- * The "offset array" of type @[n]i64@, indicating for each segment where+    --   its values begin in the data array.+    --+    -- * Finally one array per value result of the lambda, in the same order: a+    --   nonuniform result becomes its concatenation across all iterations, of+    --   outer size @[m]@, and a uniform result the collection of its values, of+    --   outer size @[n]@.+    FlatMap SubExp [VName] (ExtLambda rep)   deriving (Eq, Ord, Show)  -- | Information about computing a single histogram.@@ -381,6 +411,7 @@ data SOACMapper frep trep m = SOACMapper   { mapOnSOACSubExp :: SubExp -> m SubExp,     mapOnSOACLambda :: Lambda frep -> m (Lambda trep),+    mapOnSOACExtLambda :: ExtLambda frep -> m (ExtLambda trep),     mapOnSOACVName :: VName -> m VName   } @@ -390,6 +421,7 @@   SOACMapper     { mapOnSOACSubExp = pure,       mapOnSOACLambda = pure,+      mapOnSOACExtLambda = pure,       mapOnSOACVName = pure     } @@ -449,6 +481,11 @@             <*> mapM (mapOnSOACReduce tv) reds             <*> mapOnSOACLambda tv post_lam         )+mapSOACM tv (FlatMap w arrs lam) =+  FlatMap+    <$> mapOnSOACSubExp tv w+    <*> mapM (mapOnSOACVName tv) arrs+    <*> mapOnSOACExtLambda tv lam  mapOnSOACScan :: (Monad m) => SOACMapper frep trep m -> Scan frep -> m (Scan trep) mapOnSOACScan tv (Scan red_lam red_nes) =@@ -466,7 +503,11 @@ traverseSOACStms :: (Monad m) => OpStmsTraverser m (SOAC rep) rep traverseSOACStms f = mapSOACM mapper   where-    mapper = identitySOACMapper {mapOnSOACLambda = traverseLambdaStms f}+    mapper =+      identitySOACMapper+        { mapOnSOACLambda = traverseLambdaStms f,+          mapOnSOACExtLambda = traverseLambdaStms f+        }  instance (ASTRep rep) => FreeIn (Scan rep) where   freeIn' (Scan lam ne) = freeIn' lam <> freeIn' ne@@ -490,6 +531,7 @@         SOACMapper           { mapOnSOACSubExp = walk freeIn',             mapOnSOACLambda = walk freeIn',+            mapOnSOACExtLambda = walk freeIn',             mapOnSOACVName = walk freeIn'           } @@ -501,36 +543,99 @@         SOACMapper           { mapOnSOACSubExp = pure . substituteNames subst,             mapOnSOACLambda = pure . substituteNames subst,+            mapOnSOACExtLambda = pure . substituteNames subst,             mapOnSOACVName = pure . substituteNames subst           }  instance (ASTRep rep) => Rename (SOAC rep) where   rename = mapSOACM renamer     where-      renamer = SOACMapper rename rename rename+      renamer = SOACMapper rename rename rename rename  -- | The type of a SOAC.-soacType :: (Typed (LParamInfo rep)) => SOAC rep -> [Type]+soacType :: (Typed (LParamInfo rep)) => SOAC rep -> [ExtType] soacType (JVP shape _ _ lam) =-  lambdaReturnType lam ++ map (`arrayOfShape` shape) (lambdaReturnType lam)+  staticShapes $+    lambdaReturnType lam+      ++ map (`arrayOfShape` shape) (lambdaReturnType lam) soacType (VJP shape _ _ lam) =-  lambdaReturnType lam ++ map ((`arrayOfShape` shape) . paramType) (lambdaParams lam)+  staticShapes $ lambdaReturnType lam ++ map ((`arrayOfShape` shape) . paramType) (lambdaParams lam) soacType (WithVJP _ lam _) =-  lambdaReturnType lam+  staticShapes $ lambdaReturnType lam soacType (Stream outersize _ accs lam) =-  map (substNamesInType substs) rtp+  staticShapes $ map (substNamesInType substs) rtp   where     nms = map paramName $ take (1 + length accs) params     substs = M.fromList $ zip nms (outersize : accs)     Lambda params rtp _ = lam-soacType (Hist _ _ ops _bucket_fun) = do+soacType (Hist _ _ ops _bucket_fun) = staticShapes $ do   op <- ops   map (`arrayOfShape` histShape op) (lambdaReturnType $ histOp op) soacType (Screma w _arrs form) =-  scremaType w form+  staticShapes $ scremaType w form+soacType (FlatMap w _ lam) =+  [ Prim int64,+    arrayOfRow (Prim int64) (Free w),+    arrayOfRow (Prim Bool) (Ext 0),+    arrayOfRow (Prim int64) (Free w)+  ]+    ++ map onValueRet (drop 1 (lambdaReturnType lam))+  where+    onValueRet t+      -- An irregular result is concatenated, so its outer size becomes the+      -- total size @m@.+      | flatMapNonuniform t = t `setOuterSize` Ext 0+      -- A regular result is merely collected, one per input element.+      | otherwise = t `arrayOfRow` Free w +-- | Is this the type of an irregular result of a 'FlatMap' lambda - one that is+-- concatenated with the results of the other iterations, rather than merely+-- collected? These are exactly the results whose outermost size is the+-- existential size returned by the lambda.+flatMapNonuniform :: ExtType -> Bool+flatMapNonuniform t =+  case shapeDims $ arrayShape t of+    Ext 0 : _ -> True+    _ -> False++-- | The element types of the irregular (concatenated) arrays produced by a+-- 'FlatMap' lambda, with the existential outer dimension stripped. The+-- remaining dimensions are never existential.+flatMapRowTypes :: ExtLambda rep -> [Type]+flatMapRowTypes lam =+  fromMaybe (error "flatMapRowTypes: existential inner dimension.") $+    mapM (hasStaticShape . rowType) $+      filter flatMapNonuniform $+        drop 1 (lambdaReturnType lam)++-- | Split a list with one element per result of a 'FlatMap' into its metadata+-- results - the size, the shape array, the flag array, and the offset array -+-- and its value results.+flatMapSplitMeta :: [a] -> ((a, a, a, a), [a])+flatMapSplitMeta (m : shape : flags : offset : vals) =+  ((m, shape, flags, offset), vals)+flatMapSplitMeta _ =+  error "flatMapSplitMeta: too few results."++-- | Split a list with one element per value result of a 'FlatMap' lambda into+-- the irregular and the regular parts.+flatMapSplitValues :: ExtLambda rep -> [a] -> ([a], [a])+flatMapSplitValues lam =+  partitionEithers . zipWith f (drop 1 (lambdaReturnType lam))+  where+    f t x = if flatMapNonuniform t then Left x else Right x++-- | The types of the regular results of a 'FlatMap' lambda - the ones that are+-- collected rather than concatenated. These never have existential sizes.+flatMapUniformTypes :: ExtLambda rep -> [Type]+flatMapUniformTypes lam =+  fromMaybe (error "flatMapUniformTypes: existential size.") $+    mapM hasStaticShape $+      filter (not . flatMapNonuniform) $+        drop 1 (lambdaReturnType lam)+ instance TypedOp SOAC where-  opType = pure . staticShapes . soacType+  opType = pure . soacType  instance AliasedOp SOAC where   opAliases = map (const mempty) . soacType@@ -538,6 +643,7 @@   consumedInOp JVP {} = mempty   consumedInOp VJP {} = mempty   consumedInOp WithVJP {} = mempty+  consumedInOp FlatMap {} = mempty   -- Only map functions can consume anything.  The operands to scan   -- and reduce functions are always considered "fresh".   consumedInOp (Screma _ arrs (ScremaForm map_lam _ _ _)) =@@ -572,6 +678,8 @@       args       (Alias.analyseLambda aliases lam)       (Alias.analyseLambda aliases lam_adj)+  addOpAliases aliases (FlatMap w arrs lam) =+    FlatMap w arrs $ Alias.analyseLambda aliases lam   addOpAliases aliases (Stream size arr accs lam) =     Stream size arr accs $ Alias.analyseLambda aliases lam   addOpAliases aliases (Hist w arrs ops bucket_fun) =@@ -637,6 +745,11 @@       lam       (map depsOf' args)       <> map (const $ freeIn args <> freeIn lam) (lambdaParams lam)+  opDependencies (FlatMap w arrs lam) =+    -- FIXME: this is strictly speaking not accurate, although it will not cause+    -- trouble anytime soon.+    replicate 4 mempty+      ++ drop 1 (lambdaDependencies mempty lam (depsOfArrays w arrs))   opDependencies (Screma w arrs (ScremaForm map_lam scans reds post_lam)) =     let (scans_in, reds_in, map_deps) =           splitAt3 (scanResults scans) (redResults reds) $@@ -666,7 +779,7 @@   M.findWithDefault (Var idd) idd subs  instance CanBeWise SOAC where-  addOpWisdom = runIdentity . mapSOACM (SOACMapper pure (pure . informLambda) pure)+  addOpWisdom = runIdentity . mapSOACM (SOACMapper pure (pure . informLambda) (pure . informLambda) pure)  instance (RepTypes rep) => ST.IndexOp (SOAC rep) where   indexOp vtable k soac [i] = do@@ -747,6 +860,10 @@         </> PP.indent 2 (pretty $ lambdaReturnType lam_adj)         </> "does not match type of arguments"         </> PP.indent 2 (pretty $ map TC.argType args')+typeCheckSOAC (FlatMap w arrs lam) = do+  TC.require (Prim int64) w+  arrs' <- TC.checkSOACArrayArgs w arrs+  TC.checkExtLambda lam arrs' typeCheckSOAC (Stream size arrexps accexps lam) = do   TC.require (Prim int64) size   accargs <- mapM TC.checkArg accexps@@ -864,6 +981,8 @@     JVP shape args vec <$> rephraseLambda r lam   rephraseInOp r (WithVJP args lam lam_adj) =     WithVJP args <$> rephraseLambda r lam <*> rephraseLambda r lam_adj+  rephraseInOp r (FlatMap w arrs lam) =+    FlatMap w arrs <$> rephraseLambda r lam   rephraseInOp r (Stream w arrs acc lam) =     Stream w arrs acc <$> rephraseLambda r lam   rephraseInOp r (Hist w arrs ops lam) =@@ -896,6 +1015,8 @@   opMetrics (WithVJP _ lam lam_adj) = do     inside "WithVJP" $ lambdaMetrics lam     inside "WithVJP" $ lambdaMetrics lam_adj+  opMetrics (FlatMap _ _ lam) = do+    inside "FlatMap" $ lambdaMetrics lam   opMetrics (Stream _ _ _ lam) =     inside "Stream" $ lambdaMetrics lam   opMetrics (Hist _ _ ops bucket_fun) =@@ -933,6 +1054,13 @@             PP.braces (commasep $ map pretty args)               <> comma </> pretty lam               <> comma </> pretty lam_adj+        )+  pretty (FlatMap w arrs lam) =+    "flatmap"+      <> (parens . align)+        ( pretty w+            <> comma </> ppTuple' (map pretty arrs)+            <> comma </> pretty lam         )   pretty (Stream size arrs acc lam) =     ppStream size arrs acc lam
src/Futhark/IR/SOACS/Simplify.hs view
@@ -108,6 +108,11 @@   (lam', hoisted) <- Engine.simplifyLambda mempty lam   (lam_adj', hoisted_adj) <- Engine.simplifyLambda mempty lam_adj   pure (WithVJP args' lam' lam_adj', hoisted <> hoisted_adj)+simplifySOAC (FlatMap w arrs lam) = do+  w' <- Engine.simplify w+  arrs' <- mapM Engine.simplify arrs+  (lam', hoisted_lam) <- Engine.enterLoop $ Engine.simplifyLambda mempty lam+  pure (FlatMap w' arrs' lam', hoisted_lam) simplifySOAC (Stream outerdim arr nes lam) = do   outerdim' <- Engine.simplify outerdim   nes' <- mapM Engine.simplify nes@@ -145,7 +150,7 @@   (Simplify.SimplifiableRep rep) =>   Simplify.SimplifyOp rep (Scan (Wise rep)) simplifyScan (Scan lam nes) = do-  (lam', hoisted) <- Engine.simplifyLambda mempty lam+  (lam', hoisted) <- Engine.enterLoop $ Engine.simplifyLambda mempty lam   nes' <- Engine.simplify nes   pure (Scan lam' nes', hoisted) @@ -153,7 +158,7 @@   (Simplify.SimplifiableRep rep) =>   Simplify.SimplifyOp rep (Reduce (Wise rep)) simplifyReduce (Reduce comm lam nes) = do-  (lam', hoisted) <- Engine.simplifyLambda mempty lam+  (lam', hoisted) <- Engine.enterLoop $ Engine.simplifyLambda mempty lam   nes' <- Engine.simplify nes   pure (Reduce comm lam' nes', hoisted)
src/Futhark/IR/SegOp.hs view
@@ -962,12 +962,18 @@   pure $     Engine.hasFree bound_here       `Engine.orIf` Engine.isOp+      `Engine.orIf` isSlice       `Engine.orIf` par_blocker       `Engine.orIf` Engine.isConsumed       `Engine.orIf` Engine.isConsuming       `Engine.orIf` Engine.isDeviceMigrated   where     bound_here = namesFromList $ M.keys $ scopeOfSegSpace space+    -- Hoisting out slicing is useless and potentially dangerous, as protecting+    -- it can cause nonuniform parallelism.+    isSlice _ _ (Let _ _ (BasicOp (Index _ slice))) =+      sliceShape slice /= mempty+    isSlice _ _ _ = False  -- | We are willing to hoist potentially unsafe statements out of segops, but -- they must be protected by adding a branch on top of them.@@ -1037,8 +1043,9 @@   Engine.SimpleM rep (SegBinOp (Wise rep), Stms (Wise rep)) simplifySegBinOp phys_id (SegBinOp comm lam nes shape) = do   (lam', hoisted) <--    Engine.localVtable (\vtable -> vtable {ST.simplifyMemory = True}) $-      simplifyLambda (oneName phys_id) lam+    Engine.enterLoop+      . Engine.localVtable (\vtable -> vtable {ST.simplifyMemory = True})+      $ simplifyLambda (oneName phys_id) lam   shape' <- Engine.simplify shape   nes' <- mapM Engine.simplify nes   pure (SegBinOp comm lam' nes' shape', hoisted)@@ -1050,8 +1057,9 @@   Engine.SimpleM rep (SegPostOp (Wise rep), Stms (Wise rep)) simplifySegPostOp space (SegPostOp lam) = do   (lam', hoisted) <--    Engine.localVtable (\vtable -> vtable {ST.simplifyMemory = True}) $-      simplifyLambda bound_here lam+    Engine.enterLoop+      . Engine.localVtable (\vtable -> vtable {ST.simplifyMemory = True})+      $ simplifyLambda bound_here lam   pure (SegPostOp lam', hoisted)   where     bound_here = namesFromList $ M.keys $ scopeOfSegSpace space
src/Futhark/IR/Syntax.hs view
@@ -142,7 +142,9 @@     MatchDec (..),     MatchSort (..),     Safety (..),-    Lambda (..),+    GLambda (..),+    Lambda,+    ExtLambda,     RetAls (..),      -- * Definitions@@ -556,18 +558,26 @@     MatchEquiv   deriving (Eq, Show, Ord) --- | Anonymous function for use in a SOAC.-data Lambda rep = Lambda+-- | A generalised lambda consists of parameters, some kind of return type, and+-- a body.+data GLambda rep t = Lambda   { lambdaParams :: [LParam rep],-    lambdaReturnType :: [Type],+    lambdaReturnType :: [t],     lambdaBody :: Body rep   } -deriving instance (RepTypes rep) => Eq (Lambda rep)+deriving instance (RepTypes rep, Eq t) => Eq (GLambda rep t) -deriving instance (RepTypes rep) => Show (Lambda rep)+deriving instance (RepTypes rep, Show t) => Show (GLambda rep t) -deriving instance (RepTypes rep) => Ord (Lambda rep)+deriving instance (RepTypes rep, Ord t) => Ord (GLambda rep t)++-- | Anonymous function for use in most SOACs, where the return type is a normal+-- type.+type Lambda rep = GLambda rep Type++-- | Anonymous function that can return arrays with existential size.+type ExtLambda rep = GLambda rep ExtType  -- | A function and loop parameter. type FParam rep = Param (FParamInfo rep)
src/Futhark/IR/Syntax/Core.hs view
@@ -82,12 +82,10 @@  import Control.Category import Control.Monad-import Control.Monad.State import Data.Bifoldable import Data.Bifunctor import Data.Bitraversable import Data.Loc (locEnd, locStart)-import Data.Map.Strict qualified as M import Data.Maybe import Data.Set qualified as S import Data.String@@ -178,44 +176,13 @@   deriving (Show, Eq, Ord)  -- | A class encompassing types containing array shape information.-class (Monoid a, Eq a, Ord a) => ArrayShape a where+class (Monoid a) => ArrayShape a where   -- | Return the rank of an array with the given size.   shapeRank :: a -> Int -  -- | Check whether one shape if a subset of another shape.-  subShapeOf :: a -> a -> Bool--  -- | Prepend the dimensions of a 'Shape'.-  prependShape :: Shape -> a -> a--instance ArrayShape (ShapeBase SubExp) where+instance ArrayShape (ShapeBase a) where   shapeRank (Shape l) = length l-  subShapeOf = (==)-  prependShape = (<>) -instance ArrayShape (ShapeBase ExtSize) where-  shapeRank (Shape l) = length l-  subShapeOf (Shape ds1) (Shape ds2) =-    -- Must agree on Free dimensions, and ds1 may not be existential-    -- where ds2 is Free.  Existentials must also be congruent.-    length ds1 == length ds2-      && evalState (and <$> zipWithM subDimOf ds1 ds2) M.empty-    where-      subDimOf (Free se1) (Free se2) = pure $ se1 == se2-      subDimOf (Ext _) (Free _) = pure False-      subDimOf (Free _) (Ext _) = pure True-      subDimOf (Ext x) (Ext y) = do-        extmap <- get-        case M.lookup y extmap of-          Just ywas-            | ywas == x -> pure True-            | otherwise -> pure False-          Nothing -> do-            put $ M.insert y x extmap-            pure True--  prependShape shape = (fmap Free shape <>)- instance Semigroup Rank where   Rank x <> Rank y = Rank $ x + y @@ -224,8 +191,6 @@  instance ArrayShape Rank where   shapeRank (Rank x) = x-  subShapeOf = (==)-  prependShape shape (Rank x) = Rank $ shapeRank shape + x  -- | The memory space of a block.  If 'DefaultSpace', this is the "default" -- space, whatever that is.  The exact meaning of the 'SpaceId'@@ -303,10 +268,6 @@   | -- | Only observes value in this position, does     -- not consume.  A result may alias this.     Observe-  | -- | As 'Observe', but the result will not-    -- alias, because the parameter does not carry-    -- aliases.-    ObservePrim   deriving (Eq, Ord, Show)  -- | An identifier consists of its name and the type of the value
src/Futhark/IR/Traversals.hs view
@@ -367,6 +367,6 @@   traverseOpStms :: (Monad m) => OpStmsTraverser m (Op rep) rep  -- | A helper for defining 'traverseOpStms'.-traverseLambdaStms :: (Monad m) => OpStmsTraverser m (Lambda rep) rep+traverseLambdaStms :: (Monad m) => OpStmsTraverser m (GLambda rep t) rep traverseLambdaStms f (Lambda ps ret (Body dec stms res)) =   Lambda ps ret <$> (Body dec <$> f (scopeOfLParams ps) stms <*> pure res)
src/Futhark/IR/TypeCheck.hs view
@@ -36,6 +36,7 @@     checkArg,     checkSOACArrayArgs,     checkLambda,+    checkExtLambda,     checkBody,     consume,     binding,@@ -58,7 +59,7 @@ import Futhark.Construct (instantiateShapes) import Futhark.IR.Aliases hiding (lookupAliases) import Futhark.Util-import Futhark.Util.Pretty (align, docText, indent, ppTuple', pretty, (<+>), (</>))+import Futhark.Util.Pretty hiding (width)  -- | Information about an error during type checking.  The 'Show' -- instance for this type produces a human-readable description.@@ -736,7 +737,7 @@   lookupType ident  checkCerts :: (Checkable rep) => Certs -> TypeM rep ()-checkCerts (Certs cs) = mapM_ (requireI (Prim Unit)) cs+checkCerts = mapM_ lookupType . unCerts  checkSubExpRes :: (Checkable rep) => SubExpRes -> TypeM rep Type checkSubExpRes (SubExpRes cs se) = do@@ -758,10 +759,7 @@     delve [] =       m -checkResult ::-  (Checkable rep) =>-  Result ->-  TypeM rep ()+checkResult :: (Checkable rep) => Result -> TypeM rep () checkResult = mapM_ checkSubExpRes  checkFunBody ::@@ -777,41 +775,78 @@       matchReturnType (map fst rt) res     mapM (subExpAliasesM . resSubExp) res -checkLambdaBody ::-  (Checkable rep) =>-  [Type] ->-  Body (Aliases rep) ->-  TypeM rep ()-checkLambdaBody ret (Body (_, rep) stms res) = do-  checkBodyDec rep-  checkStms stms $ checkLambdaResult ret res--checkLambdaResult ::-  (Checkable rep) =>-  [Type] ->-  Result ->-  TypeM rep ()-checkLambdaResult ts es-  | length ts /= length es =+checkNumResults :: (Pretty t) => [t] -> Result -> TypeM rep ()+checkNumResults ts res+  | length ts /= length res =       bad . TypeError $         "Lambda has return type "           <> prettyTuple ts           <> " describing "           <> prettyText (length ts)           <> " values, but body returns "-          <> prettyText (length es)+          <> prettyText (length res)           <> " values: "-          <> prettyTuple es-  | otherwise = forM_ (zip ts es) $ \(t, e) -> do-      et <- checkSubExpRes e-      unless (et == t) . bad . TypeError $-        "Subexpression "-          <> prettyText e-          <> " has type "-          <> prettyText et-          <> " but expected "-          <> prettyText t+          <> prettyTuple res+  | otherwise = pure () +checkLambdaResult :: (Checkable rep) => [Type] -> Result -> TypeM rep ()+checkLambdaResult ts es = do+  checkNumResults ts es+  forM_ (zip ts es) $ \(t, e) -> do+    et <- checkSubExpRes e+    unless (et == t) . bad . TypeError $+      "Subexpression "+        <> prettyText e+        <> " has type "+        <> prettyText et+        <> " but expected "+        <> prettyText t++checkLambdaBody ::+  (Checkable rep) => [Type] -> Body (Aliases rep) -> TypeM rep ()+checkLambdaBody ret (Body (_, rep) stms res) = do+  mapM_ checkType ret+  checkBodyDec rep+  checkStms stms $ do+    checkLambdaResult ret res++checkExtLambdaResult :: (Checkable rep) => [ExtType] -> Result -> TypeM rep ()+checkExtLambdaResult ts es = do+  checkNumResults ts es+  matchExtReturnType ts es++checkExtLambdaBody ::+  (Checkable rep) => [ExtType] -> Body (Aliases rep) -> TypeM rep ()+checkExtLambdaBody ret (Body (_, rep) stms res) = do+  mapM_ checkExtType ret+  case ret of+    Prim (IntType Int64) : ts -> mapM_ checkValueRet ts+    _ ->+      bad . TypeError $+        "flatmap lambda must return a size of type i64 first, but returns "+          <> prettyTuple ret+  checkBodyDec rep+  checkStms stms $ checkExtLambdaResult ret res+  where+    -- A result is either a segment, in which case the existential size is its+    -- outermost size, or a value, in which case it has no existential size at+    -- all.+    checkValueRet t =+      case shapeDims $ arrayShape t of+        _ : inner+          | any (isJust . isExt) inner ->+              bad . TypeError $+                "flatmap result type "+                  <> prettyText t+                  <> " has existential in inner dimension."+        Ext i : _+          | i /= 0 ->+              bad . TypeError $+                "flatmap result type "+                  <> prettyText t+                  <> " has an existential size other than the first result."+        _ -> pure ()+ checkBody ::   (Checkable rep) =>   Body (Aliases rep) ->@@ -986,7 +1021,7 @@           rettype_ext           (staticShapes bodyt)     Just rettype' ->-      unless (bodyt `subtypesOf` rettype') . bad $+      unless (bodyt == rettype') . bad $         ReturnTypeError           (nameFromString "<loop body>")           (staticShapes rettype')@@ -1026,9 +1061,9 @@   when (rettype_derived /= rettype_annot) $     bad . TypeError . docText $       "Expected apply result type:"-        </> indent 2 (pretty $ map fst rettype_derived)+        </> indent 2 (braces $ commasep $ map prettyRet rettype_derived)         </> "But annotation is:"-        </> indent 2 (pretty $ map fst rettype_annot)+        </> indent 2 (braces $ commasep $ map prettyRet rettype_annot)   consumeArgs paramtypes argflows checkExp (Loop merge form loopbody) = do   let mergepat = map fst merge@@ -1134,7 +1169,8 @@      pure (Acc (paramName p) shape elem_ts NoUniqueness, mempty) -  checkAnyLambda False lam $ replicate num_accs (Prim Unit, mempty) ++ acc_args+  checkAnyLambda False checkLambdaBody lam $+    replicate num_accs (Prim Unit, mempty) ++ acc_args   where     num_accs = length inputs checkExp (Op op) = do@@ -1242,9 +1278,8 @@   TypeM rep a ->   TypeM rep a checkStm stm@(Let pat aux e) m = do-  let Certs cs = stmAuxCerts aux-      (_, dec) = stmAuxDec aux-  context "When checking certificates" $ mapM_ (requireI $ Prim Unit) cs+  let (_, dec) = stmAuxDec aux+  context "When checking certificates" $ checkCerts $ stmAuxCerts aux   context "When checking expression annotation" $ checkExpDec dec   context ("When matching\n" <> message "  " pat <> "\nwith\n" <> message "  " e) $     matchPat pat e@@ -1352,8 +1387,13 @@ -- The boolean indicates whether we only allow consumption of -- parameters. checkAnyLambda ::-  (Checkable rep) => Bool -> Lambda (Aliases rep) -> [Arg] -> TypeM rep ()-checkAnyLambda soac (Lambda params rettype body) args = do+  (Checkable rep) =>+  Bool ->+  ([t] -> Body (Aliases rep) -> TypeM rep ()) ->+  GLambda (Aliases rep) t ->+  [Arg] ->+  TypeM rep ()+checkAnyLambda soac onBody (Lambda params rettype body) args = do   let fname = nameFromString "<anonymous>"   if length params == length args     then do@@ -1372,8 +1412,7 @@       binding (M.fromList params') $         maybe id consumeOnlyParams consumable $ do           checkLambdaParams params-          mapM_ checkType rettype-          checkLambdaBody rettype body+          onBody rettype body     else       bad . TypeError $         "Anonymous function defined with "@@ -1385,7 +1424,10 @@           <> " arguments."  checkLambda :: (Checkable rep) => Lambda (Aliases rep) -> [Arg] -> TypeM rep ()-checkLambda = checkAnyLambda True+checkLambda = checkAnyLambda True checkLambdaBody++checkExtLambda :: (Checkable rep) => ExtLambda (Aliases rep) -> [Arg] -> TypeM rep ()+checkExtLambda = checkAnyLambda True checkExtLambdaBody  checkPrimExp :: (Checkable rep) => PrimExp VName -> TypeM rep () checkPrimExp ValueExp {} = pure ()
src/Futhark/Internalise/Bindings.hs view
@@ -4,6 +4,9 @@ module Futhark.Internalise.Bindings   ( internaliseAttrs,     internaliseAttr,+    FunParams (..),+    internaliseFParams,+    bindFParams,     bindingFParams,     bindingLoopParams,     bindingLambdaParams,@@ -40,32 +43,53 @@ treeLike (Pure _) _ = error "treeLike: invalid input" treeLike (Free ls) bs = Free $ zipWith treeLike ls (chunks (map length ls) bs) -bindingFParams ::+-- | The internalised parameters of a function. Computing these assigns fresh+-- names, so we split it from binding them in scope ('bindFParams'): that way a+-- caller can internalise the parameters once, and both register the function's+-- calling information and later bind the parameters for its body without+-- recomputing them.+data FunParams = FunParams+  { -- | Shape (and certificate) parameters.+    funShapeParams :: [I.FParam I.SOACS],+    -- | Value parameters, grouped as the source parameters.+    funValueParams :: [[Tree (I.FParam I.SOACS)]],+    funSubsts :: VarSubsts,+    funShapeSubst :: VarSubsts+  }++-- | All parameters, flattened; used for scoping and the 'I.FunDef'.+funAllParams :: FunParams -> [I.FParam I.SOACS]+funAllParams fps =+  funShapeParams fps ++ foldMap (foldMap toList) (funValueParams fps)++-- | Internalise a function's parameters (see 'FunParams').+internaliseFParams ::   [E.TypeParam] ->   [E.Pat E.ParamType] ->-  ([I.FParam I.SOACS] -> [[Tree (I.FParam I.SOACS)]] -> InternaliseM a) ->-  InternaliseM a-bindingFParams tparams params m = do+  InternaliseM FunParams+internaliseFParams tparams params = do   flattened_params <- mapM flattenPat params   let params_idents = concat flattened_params   params_ts <-     internaliseParamTypes $       map (E.unInfo . E.identType . fst) params_idents   let num_param_idents = map length flattened_params--  let shape_params = [I.Param mempty v $ I.Prim I.int64 | E.TypeParamDim v _ <- tparams]+      shape_params = [I.Param mempty v $ I.Prim I.int64 | E.TypeParamDim v _ <- tparams]       shape_subst = M.fromList [(I.paramName p, [I.Var $ I.paramName p]) | p <- shape_params]-  bindingFlatPat params_idents (concatMap (concatMap toList) params_ts) $ \valueparams -> do-    let (certparams, valueparams') =-          first concat $ unzip $ map fixAccParams valueparams-        all_params = certparams ++ shape_params ++ concat valueparams'-    I.localScope (I.scopeOfFParams all_params) $-      substitutingVars shape_subst $ do-        let values_grouped_by_params = chunks num_param_idents valueparams'-            types_grouped_by_params = chunks num_param_idents params_ts--        m (certparams ++ shape_params) $-          zipWith chunkValues types_grouped_by_params values_grouped_by_params+  (valueparams, substs) <-+    processFlatPat params_idents (concatMap (concatMap toList) params_ts)+  let (certparams, valueparams') =+        first concat $ unzip $ map fixAccParams valueparams+      values_grouped_by_params = chunks num_param_idents valueparams'+      types_grouped_by_params = chunks num_param_idents params_ts+  pure+    FunParams+      { funShapeParams = certparams ++ shape_params,+        funValueParams =+          zipWith chunkValues types_grouped_by_params values_grouped_by_params,+        funSubsts = substs,+        funShapeSubst = shape_subst+      }   where     fixAccParams ps =       first catMaybes $ unzip $ map fixAccParam ps@@ -83,6 +107,22 @@       concat $ zipWith f tss vss       where         f ts vs = zipWith treeLike ts (chunks (map length ts) vs)++-- | Bind already-internalised parameters (see 'internaliseFParams') in scope.+bindFParams :: FunParams -> InternaliseM a -> InternaliseM a+bindFParams fps m =+  local (\env -> env {envSubsts = funSubsts fps `M.union` envSubsts env}) $+    I.localScope (I.scopeOfFParams (funAllParams fps)) $+      substitutingVars (funShapeSubst fps) m++bindingFParams ::+  [E.TypeParam] ->+  [E.Pat E.ParamType] ->+  ([I.FParam I.SOACS] -> [[Tree (I.FParam I.SOACS)]] -> InternaliseM a) ->+  InternaliseM a+bindingFParams tparams params m = do+  fps <- internaliseFParams tparams params+  bindFParams fps $ m (funShapeParams fps) (funValueParams fps)  bindingLoopParams ::   [E.TypeParam] ->
src/Futhark/Internalise/Defunctionalise.hs view
@@ -22,6 +22,9 @@  -- | A static value stores additional information about the result of -- defunctionalization of an expression, aside from the residual expression.+--+-- The Ord instance here is really important, as it is used for the memoisation+-- machinery that handles recursive functions. data StaticVal   = Dynamic ParamType   | -- | The Env is the lexical closure of the lambda.@@ -37,7 +40,7 @@     DynamicFun (Exp, StaticVal) StaticVal   | IntrinsicSV   | HoleSV StructType SrcLoc-  deriving (Show)+  deriving (Show, Eq, Ord)  data Binding = Binding   { -- | Just if this is a polymorphic binding that must be@@ -45,7 +48,7 @@     bindingType :: Maybe ([VName], StructType),     bindingSV :: StaticVal   }-  deriving (Show)+  deriving (Show, Eq, Ord)  -- | Environment mapping variable names to their associated static -- value.@@ -175,33 +178,51 @@     restrict' (HoleSV t loc) = HoleSV t loc     restrict'' (Binding t sv) = Binding t $ restrict' sv +-- | Maps a function specialisation - a function's 'LambdaSV' together with the+-- 'StaticVal' of the argument it is applied to - to the top-level function+-- lifted for it: its name, its return type, and the 'StaticVal' of the result.+-- Applying the same function to the same static argument again (in particular, a+-- recursive occurrence) reuses this lifting rather than lifting it anew, which+-- both deduplicates lifted functions and ties off recursion. See Note [Lifting+-- and recursion].+type LiftMemo = M.Map (StaticVal, StaticVal) (VName, ResRetType, StaticVal)+ -- | Defunctionalization monad. The Reader environment tracks both the global -- Env and the current Env. This is used to avoid unnecessarily large closure -- environments (no need to capture the global one). newtype DefM a-  = DefM (ReaderT (Env, Env) (State ([ValBind], VNameSource)) a)+  = DefM (ReaderT (Env, Env) (State ([ValBind], VNameSource, LiftMemo)) a)   deriving     ( Functor,       Applicative,       Monad,       MonadReader (Env, Env),-      MonadState ([ValBind], VNameSource)+      MonadState ([ValBind], VNameSource, LiftMemo)     )  instance MonadFreshNames DefM where-  putNameSource src = modify $ \(x, _) -> (x, src)-  getNameSource = gets snd+  putNameSource src = modify $ \(x, _, m) -> (x, src, m)+  getNameSource = gets (\(_, src, _) -> src)  -- | Run a computation in the defunctionalization monad. Returns the result of -- the computation, a new name source, and a list of lifted function declations. runDefM :: VNameSource -> DefM a -> (a, VNameSource, [ValBind]) runDefM src (DefM m) =-  let (x, (vbs, src')) = runState (runReaderT m mempty) (mempty, src)+  let (x, (vbs, src', _)) = runState (runReaderT m mempty) (mempty, src, mempty)    in (x, src', reverse vbs)  addValBind :: ValBind -> DefM ()-addValBind vb = modify $ first (vb :)+addValBind vb = modify $ \(vbs, src, m) -> (vb : vbs, src, m) +-- | The function lifted for the given specialisation, if any (see 'LiftMemo').+lookupLift :: (StaticVal, StaticVal) -> DefM (Maybe (VName, ResRetType, StaticVal))+lookupLift key = gets $ \(_, _, m) -> M.lookup key m++-- | Record the function lifted for a specialisation (see 'LiftMemo').+insertLift :: (StaticVal, StaticVal) -> (VName, ResRetType, StaticVal) -> DefM ()+insertLift key v =+  modify $ \(vbs, src, m) -> (vbs, src, M.insert key v m)+ -- | Create a new top-level value declaration with the given function name, -- return type, list of parameters, and body expression. liftValDec :: VName -> ResRetType -> [VName] -> [Pat ParamType] -> Exp -> DefM ()@@ -865,7 +886,7 @@  defuncApplyFunction :: Exp -> Int -> DefM (Exp, StaticVal) defuncApplyFunction e@(Var qn (Info t) loc) num_args = do-  let (argtypes, rettype) = unfoldFunType t+  let (argtypes, rettype) = first (map snd) $ unfoldFunType t   sv <- lookupVar (toStruct t) (qualLeaf qn)    case sv of@@ -920,60 +941,91 @@   (Exp, StaticVal) ->   ((Maybe VName, Exp), [ParamType]) ->   DefM (Exp, StaticVal)-defuncApplyArg (fname_s, floc) (f', LambdaSV pat lam_e_t lam_e closure_env) ((argext, arg), _) = do+defuncApplyArg (fname_s, floc) (f', fsv@(LambdaSV pat lam_e_t lam_e closure_env)) ((argext, arg), _) = do   (arg', arg_sv) <- defuncExp arg-  let env' = alwaysMatchPatSV closure_env pat arg_sv-      dims = mempty-  (lam_e', sv) <--    localNewEnv env' $-      defuncExp lam_e -  let closure_pat = buildEnvPat dims closure_env-      pat' = updatePat pat arg_sv+  -- Build a call to the lifted specialisation named 'lifted' with the+  -- given (first-order) return type.  Reuses the closure value 'f'' and+  -- the (already defunctionalised) argument.+  let mkCall lifted lifted_ret = do+        let f_t = toStruct $ typeOf f'+            arg_t = toStruct $ typeOf arg'+            fname_t = foldFunType [toParam Observe f_t, toParam (diet (patternType pat)) arg_t] lifted_ret+            fname' = Var (qualName lifted) (Info fname_t) floc+        callret <- unRetType lifted_ret+        pure $ mkApply fname' [(Nothing, f'), (argext, arg')] callret -  globals <- asks $ M.keysSet . fst+  -- Applying this function to this static argument may already have been+  -- lifted; if so, reuse that lifting. This is what ties off recursion,+  -- since a recursive occurrence re-applies the same function to the same+  -- static argument. See Note [Lifting and recursion].+  let key = (fsv, arg_sv) -  -- 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']-      lifted_rettype =-        RetType (retDims lam_e_t) $-          combineTypeShapes (retType lam_e_t) (resTypeFromSV sv)+  memhit <- lookupLift key+  case memhit of+    Just (lifted, lifted_ret, sv) -> do+      e <- mkCall lifted lifted_ret+      pure (e, sv)+    Nothing -> do+      let env' = alwaysMatchPatSV closure_env pat arg_sv+          dims = mempty -      already_bound =-        globals <> S.fromList (dims <> foldMap patNames params)+      -- This slot yields the first-order body of the function (rather+      -- than peeling off another function-typed parameter) exactly when+      -- its result is order-zero.  There we must record the lifted name+      -- before defunctionalising the body, so that a recursive+      -- occurrence inside it finds the entry rather than re-inlining+      -- forever.  The recorded return type is the declared one, since+      -- the body's inferred refinement is not available yet; we replace+      -- it with the refined one once the body has been lifted.+      let is_body = orderZero $ retType lam_e_t+      fname <- newVName fname_s+      let memo_ret = RetType (retDims lam_e_t) (retType lam_e_t)+      when is_body $+        insertLift key (fname, memo_ret, Dynamic $ resToParam $ retType memo_ret)+      (lam_e', sv) <-+        localNewEnv env' $+          defuncExp lam_e -      more_dims =-        S.toList $-          S.filter (`S.notMember` already_bound) $-            foldMap patternArraySizes params+      let closure_pat = buildEnvPat dims closure_env+          pat' = updatePat pat arg_sv -  -- 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-  params' <- instAnySizes params+      globals <- asks $ M.keysSet . fst -  fname <- newVName fname_s-  liftValDec-    fname-    lifted_rettype-    (dims ++ more_dims ++ unboundSizes bound_sizes params')-    params'-    lam_e'+      -- 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']+          lifted_rettype =+            RetType (retDims lam_e_t) $+              combineTypeShapes (retType lam_e_t) (resTypeFromSV sv) -  let f_t = toStruct $ typeOf f'-      arg_t = toStruct $ typeOf arg'-      fname_t = foldFunType [toParam Observe f_t, toParam (diet (patternType pat)) arg_t] lifted_rettype-      fname' = Var (qualName fname) (Info fname_t) floc-  callret <- unRetType lifted_rettype+          already_bound =+            globals <> S.fromList (dims <> foldMap patNames params) -  pure-    ( mkApply fname' [(Nothing, f'), (argext, arg')] callret,-      sv-    )+          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+      params' <- instAnySizes params++      liftValDec+        fname+        lifted_rettype+        (dims ++ more_dims ++ unboundSizes bound_sizes params')+        params'+        lam_e'++      insertLift key (fname, lifted_rettype, sv)++      e <- mkCall fname lifted_rettype+      pure (e, 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.@@ -1016,8 +1068,8 @@           (argtypes, _) = unfoldFunType $ typeOf f       fmap (first $ updateReturn appres) $         foldM (defuncApplyArg (fname, loc)) (f', f_sv) $-          NE.zip args $-            NE.tails argtypes+          NE.zip args . NE.tails . map snd $+            argtypes   where     intrinsicOrHole e' = do       -- If the intrinsic is fully applied, then we are done.@@ -1240,12 +1292,52 @@ svFromType (Scalar (Record fs)) = RecordSV . M.toList $ M.map svFromType fs svFromType t = Dynamic t +-- | The static value of a (top-level) binding within the scope of its own+-- definition, used to handle recursion. We construct a 'DynamicFun' spine+-- matching the value parameters, ending in a 'Dynamic' for the result. This+-- suffices for fully applied recursive calls: the residual expression simply+-- refers to the top-level binding by name.+--+-- This is possible only when the parameters and the result are order-zero, as+-- the static value would otherwise depend on the body. See Note [Lifting and+-- recursion] for why that suffices.+selfSV :: VName -> [Pat ParamType] -> ResType -> Maybe StaticVal+selfSV name params rettype+  | all patternOrderZero params,+    orderZero rettype =+      Just $ go params+  | otherwise = Nothing+  where+    ret_sv = Dynamic $ resToParam rettype+    go [] = ret_sv+    go (_ : ps) =+      let inner = go ps+          self = Var (qualName name) (Info (structTypeFromSV inner)) mempty+       in DynamicFun (self, inner) inner++-- | The self static value binding for a top-level binding, if it is first-order+-- (see 'selfSV'). We register these for all bindings before defunctionalising+-- any of them, so that mutually recursive references resolve. The real static+-- value of each binding, added as it is processed, takes precedence over the+-- bindings produced here. This is what lets a lambda-lifted function refer to+-- the binding it was lifted out of; see Note [Lifting and recursion].+selfBinding :: ValBind -> Env+selfBinding valbind =+  case selfSV (valBindName valbind) (valBindParams valbind) rettype of+    Just self_sv ->+      M.singleton (valBindName valbind) $+        Binding (Just (first (map typeParamName) (valBindTypeScheme valbind))) self_sv+    Nothing -> mempty+  where+    Info (RetType _ rettype) = valBindRetType valbind+ -- | Defunctionalize a top-level value binding. Returns the -- transformed result as well as an environment that binds the name of--- the value binding to the static value of the transformed body.  The--- boolean is true if the function is a 'DynamicFun'.+-- the value binding to the static value of the transformed body. defuncValBind :: ValBind -> DefM (ValBind, Env)--- Eta-expand entry points with a functional return type.+-- Eta-expand a functional result into further parameters, as the core language+-- is first-order. A recursive binding never has one (see Note [Lifting and+-- recursion]), so this cannot affect the self static values. defuncValBind (ValBind entry name name_loc _ (Info rettype) tparams params body _ attrs loc)   | Scalar Arrow {} <- retType rettype = do       (body_pats, body', rettype') <- etaExpand (second (const mempty) rettype) body@@ -1268,8 +1360,24 @@       show name         ++ " has type parameters, "         ++ "but the defunctionaliser expects a monomorphic input program."+  -- Bind the name in the scope of its own definition, so that first-order+  -- recursive references can be defunctionalised via a static value; see+  -- 'selfSV'. Higher-order recursive functions get no self static value: their+  -- body is not defunctionalised here (it is stored raw in a 'LambdaSV' and+  -- defunctionalised when applied), so recursive references resolve against the+  -- global scope, and the recursion is tied off by memoisation in+  -- 'defuncApplyArg' (see Note [Lifting and recursion]).+  let self = (selfBinding valbind <>)+  -- The self static value goes into the *global* environment (as well as the+  -- local one) so that it is treated as a top-level function: recursive uses of+  -- it as a value are then eta-expanded and closure-converted like any other+  -- top-level function (see the 'DynamicFun' case of 'defuncExp'), rather than+  -- capturing the self static value into a closure, which 'restrictEnvTo' omits+  -- for globals.   (tparams', params', body', sv, sv_t) <--    defuncLet (map typeParamName tparams) params body $ RetType ret_dims rettype+    local (bimap self self) $+      defuncLet (map typeParamName tparams) params body $+        RetType ret_dims rettype   globals <- asks $ M.keysSet . fst   let bound_sizes = S.fromList (foldMap patNames params') <> S.fromList tparams' <> globals   params'' <- instAnySizes params'@@ -1297,11 +1405,15 @@  -- | Defunctionalize a list of top-level declarations. defuncVals :: [ValBind] -> DefM ()-defuncVals [] = pure ()-defuncVals (valbind : ds) = do-  (valbind', env) <- defuncValBind valbind-  addValBind valbind'-  local (bimap (env <>) (env <>)) $ defuncVals ds+defuncVals binds =+  local (bimap (self_env <>) (self_env <>)) $ go binds+  where+    self_env = foldMap selfBinding binds+    go [] = pure ()+    go (valbind : ds) = do+      (valbind', env) <- defuncValBind valbind+      addValBind valbind'+      local (bimap (env <>) (env <>)) $ go ds  {-# NOINLINE transformProg #-} @@ -1312,3 +1424,74 @@ transformProg decs = modifyNameSource $ \namesrc ->   let ((), namesrc', decs') = runDefM namesrc $ defuncVals decs    in (decs', namesrc')++-- Note [Lifting and recursion]+--+-- Defunctionalisation works by static interpretation: applying a higher-order+-- function ('defuncApplyArg' on a 'LambdaSV') is specialised by inlining the+-- function body with its static (function-valued) argument substituted in, and+-- lifting the result to a fresh top-level function. Recursion is not a special+-- case of this; it falls out of memoising the lifting.+--+-- The specialisation a given step produces is fully determined by the function+-- (the 'LambdaSV', which carries the body, the parameter, and the lexical+-- closure) together with the 'StaticVal' of the argument it is applied to.+-- Everything that varies at runtime (the captured values) is passed in the+-- closure, not baked into the lifted function. So we memoise ('LiftMemo') the+-- lifted function under exactly that pair, and reuse it whenever the same+-- function is applied to the same static argument again. This deduplicates+-- lifted functions, and it terminates recursion for free: a recursive+-- occurrence re-applies the same function to a statically identical argument+-- (order-zero arguments have the same 'StaticVal' regardless of their runtime+-- value, and an unchanged higher-order argument resolves to the same+-- 'StaticVal'), so it hits the memo and becomes a call to the function being+-- lifted rather than another round of inlining.+--+-- There is no distinction between recursive and non-recursive functions: an+-- ordinary call and a recursive call are the same memo lookup. Note that this+-- means a program whose recursion is not statically resolvable - an indirect+-- recursive call, or one that passes a *different* function at each step (e.g.+-- `go (\x -> g x) n = ... go (\x -> g (g x)) (n-1)`) - produces ever-changing+-- keys and does not terminate here. The type checker rejects such programs, by+-- requiring a recursive application to be saturated and to pass its+-- higher-order arguments unchanged.+--+-- One subtlety remains, and it is about knot-tying, not about recognising+-- recursion. A curried application peels off one parameter per step; the step+-- whose result is order-zero yields the actual first-order body, and its body+-- may contain the recursive occurrence. That entry must therefore be recorded+-- _before_ the body is defunctionalised, using the declared (order-zero) return+-- type and a 'Dynamic' static value, both known up front; it is replaced with+-- the refined return type and real static value once the body is lifted. The+-- earlier (function-typed) steps are recorded *after* lifting, since their+-- residual 'StaticVal' - which a memo hit must return so application can+-- continue - is not known until then. This is sound because those earlier steps+-- are all processed before the order-zero step (the body is innermost), so they+-- are in the memo by the time the body's recursive occurrence looks for them.+--+-- Another problem is scoping. Recursion is resolved by name, and 'defuncVals'+-- processes bindings in order, so a body being defunctionalised can normally+-- see the static value of everything it refers to. Lambda lifting is the only+-- thing that breaks this as it can introduce mutual recursion (this is not+-- allowed in the source language).+--+-- Whether this is a problem depends on the lifted function:+--+-- - A higher-order function is not defunctionalised where it is defined - its+--   body is stored in a 'LambdaSV' and defunctionalised when it is applied,+--   which is necessarily from a later binding. By then the binding it refers to+--   has been processed, and nothing more is needed.+--+-- - A first-order function is defunctionalised eagerly, and so does need the+--   static value of a binding that has not been processed yet. 'selfBinding'+--   provides it for all bindings up front, which 'selfSV' can derive from the+--   signature alone when the parameters and the result are order-zero.+--+-- The type rules ensure the result of a recursive binding is always order-zero,+-- which is crucial for making this work.+--+-- A binding with a higher-order parameter thus gets no self static value, and+-- needs none: the type checker also requires a recursive application to pass+-- such a parameter unchanged, so a lifted function containing the recursive+-- occurrence must capture the function-typed parameter, which makes it+-- higher-order itself - the first case above.
src/Futhark/Internalise/Exps.hs view
@@ -44,76 +44,98 @@   I.renameProg $ I.Prog opaques consts funs  internaliseValBinds :: VisibleTypes -> [E.ValBind] -> InternaliseM ()-internaliseValBinds types = mapM_ $ internaliseValBind types+internaliseValBinds types vbinds = do+  -- Internalise every function's parameters and register its calling+  -- information *before* internalising any body, so that references between+  -- functions resolve regardless of definition order.+  headers <- forM vbinds $ \vb -> do+    let E.ValBind _ fname _ _ (Info rettype) tparams params _ _ _ _ = vb+    fps <- internaliseFParams tparams params+    let (_, _, _, info) = funHeader (funShapeParams fps) (funValueParams fps) rettype+    unless (null (funValueParams fps)) $ addFunInfo fname info+    pure (vb, fps)+  forM_ headers $ \(vb, fps) ->+    bindFParams fps $ internaliseValBindBody types vb fps +-- | A function's internalised parameters, return type, full internalised+-- return type, and calling information, all computed from its signature alone+-- (no body required).+funHeader ::+  [I.FParam I.SOACS] ->+  [[Tree (I.FParam I.SOACS)]] ->+  E.ResRetType ->+  ([Tree (I.FParam I.SOACS)], [I.DeclExtType], [(I.DeclExtType, RetAls)], FunInfo)+funHeader shapeparams params' rettype =+  (all_params, rettype', fun_rettype, info)+  where+    shapenames = map I.paramName shapeparams+    all_params = map pure shapeparams ++ concat params'+    zeroExts ts = generaliseExtTypes ts ts+    (rettype', retals) =+      first zeroExts . unzip $+        internaliseReturnType (map (fmap paramDeclType) all_params) rettype+    num_ctx = length (shapeContext rettype')+    fun_rettype =+      replicate num_ctx (I.Prim int64, mempty)+        ++ zip rettype' (map (shiftRetAls num_ctx) retals)+    info =+      ( shapenames,+        map declTypeOf $ foldMap (foldMap toList) params',+        foldMap toList all_params,+        fmap (`zip` map snd fun_rettype)+          . applyRetType (map fst fun_rettype) (foldMap toList all_params)+      )+ internaliseFunName :: VName -> Name internaliseFunName = nameFromText . prettyText  shiftRetAls :: Int -> RetAls -> RetAls shiftRetAls d (RetAls pals rals) = RetAls pals $ map (+ d) rals -internaliseValBind :: VisibleTypes -> E.ValBind -> InternaliseM ()-internaliseValBind types fb@(E.ValBind entry fname _ _ (Info rettype) tparams params body _ attrs _) = do-  bindingFParams tparams params $ \shapeparams params' -> do-    let shapenames = map I.paramName shapeparams-        all_params = map pure shapeparams ++ concat params'-        msg =-          errorMsg-            [ "Internal runtime error.\n",-              "Return value of ",-              ErrorString (prettyText fname),-              " does not match type shape.\n",-              "This is a bug in the Futhark compiler. Please report this:\n",-              "  https://github.com/diku-dk/futhark/issues"-            ]+-- | Internalise a function's body, given its already-bound parameters (see+-- 'internaliseFParams' and 'bindFParams'). Its calling information has already+-- been registered by 'internaliseValBinds'.+internaliseValBindBody :: VisibleTypes -> E.ValBind -> FunParams -> InternaliseM ()+internaliseValBindBody types fb@(E.ValBind entry fname _ _ (Info rettype) _ _ body _ attrs _) fps = do+  let (all_params, rettype', fun_rettype, _) =+        funHeader (funShapeParams fps) (funValueParams fps) rettype+      params' = funValueParams fps+      msg =+        errorMsg+          [ "Internal runtime error.\n",+            "Return value of ",+            ErrorString (prettyText fname),+            " does not match type shape.\n",+            "This is a bug in the Futhark compiler. Please report this:\n",+            "  https://github.com/diku-dk/futhark/issues"+          ] -    (body', rettype') <- buildBody $ do-      body_res <- internaliseExp (baseName fname <> "_res") body-      (rettype', retals) <--        first zeroExts . unzip . internaliseReturnType (map (fmap paramDeclType) all_params) rettype-          <$> mapM subExpType body_res+  body' <- buildBody_ $ do+    body_res <- internaliseExp (baseName fname <> "_res") body -      when (null params') $-        bindExtSizes (E.AppRes (E.toStruct $ E.retType rettype) (E.retDims rettype)) body_res+    when (null params') $+      bindExtSizes (E.AppRes (E.toStruct $ E.retType rettype) (E.retDims rettype)) body_res -      body_res' <--        ensureResultExtShape msg (map I.fromDecl rettype') $ subExpsRes body_res-      let num_ctx = length (shapeContext rettype')-      pure-        ( body_res',-          replicate num_ctx (I.Prim int64, mempty)-            ++ zip rettype' (map (shiftRetAls num_ctx) retals)-        )+    ensureResultExtShape msg (map I.fromDecl rettype') $ subExpsRes body_res -    attrs' <- internaliseAttrs attrs+  attrs' <- internaliseAttrs attrs -    let fd =-          I.FunDef-            Nothing-            attrs'-            (internaliseFunName fname)-            rettype'-            (foldMap toList all_params)-            body'+  let fd =+        I.FunDef+          Nothing+          attrs'+          (internaliseFunName fname)+          fun_rettype+          (foldMap toList all_params)+          body' -    if null params'-      then bindConstant fname fd-      else-        bindFunction-          fname-          fd-          ( shapenames,-            map declTypeOf $ foldMap (foldMap toList) params',-            foldMap toList all_params,-            fmap (`zip` map snd rettype')-              . applyRetType (map fst rettype') (foldMap toList all_params)-          )+  if null params'+    then addConstant fname fd+    else addFunDef fd    case entry of     Just (Info entry') -> generateEntryPoint types entry' fb     Nothing -> pure ()-  where-    zeroExts ts = generaliseExtTypes ts ts  generateEntryPoint :: VisibleTypes -> E.EntryPoint -> E.ValBind -> InternaliseM () generateEntryPoint types (E.EntryPoint e_params e_rettype doc) vb = do@@ -558,13 +580,17 @@             -- not in the same position), in which case we must be careful to             -- avoid clobbering.             let mergepat_names = map I.paramName mergepat'-            ses' <- forM ses $ \case-              I.Var v+            ses' <- forM (zip mergepat' ses) $ \case+              (p, I.Var v)                 | v `elem` mergepat_names -> do                     v' <- newVName $ baseName v <> "_tmp"-                    letBindNames [v'] $ I.BasicOp (I.SubExp $ I.Var v)+                    letBindNames [v'] $+                      if primType $ paramType p+                        then I.BasicOp (I.SubExp $ I.Var v)+                        -- Need administrative coerce due to the renaming.+                        else shapeCoerce (I.arrayDims $ paramType p) v                     pure $ I.Var v'-              se -> pure se+              (_, se) -> pure se             forM_ (zip mergepat' ses') $ \(p, se) ->               letBindNames [I.paramName p] $                 case se of@@ -722,7 +748,7 @@   internaliseExp desc e internaliseExp desc (E.Coerce e _ (Info et) _) = do   ses <- internaliseExp desc e-  ts <- internaliseCoerceType (E.toStruct et) <$> mapM subExpType ses+  let ts = internaliseCoerceType (E.toStruct et)   dt' <- typeExpForError $ toStruct et   forM (zip ses ts) $ \(e', t') -> do     dims <- arrayDims <$> subExpType e'@@ -909,8 +935,8 @@         case t of           I.Array pt shape _ ->             letSubExp desc $ I.BasicOp $ I.Scratch pt $ I.shapeDims shape-          I.Prim pt ->-            pure $ constant $ blankPrimValue pt+          -- Ignore scratch on non-arrays because they are sometimes applied+          -- too promisciously.           _ -> pure se     "blank" -> do       ts <- mapM subExpType e'@@ -1456,12 +1482,12 @@     E.Scalar (E.Prim (E.Signed it)) -> (,it) <$> asIntS Int64 e'     _ -> error "internaliseSizeExp: bad type" +asVar :: Name -> I.SubExp -> InternaliseM I.VName+asVar desc se = letExp desc $ I.BasicOp $ I.SubExp se+ internaliseExpToVars :: Name -> E.Exp -> InternaliseM [I.VName] internaliseExpToVars desc e =-  mapM asIdent =<< internaliseExp desc e-  where-    asIdent (I.Var v) = pure v-    asIdent se = letExp desc $ I.BasicOp $ I.SubExp se+  mapM (asVar desc) =<< internaliseExp desc e  internaliseOperation ::   Name ->@@ -1696,6 +1722,55 @@       rettype       =<< bodyBind body +-- | The number of internalised values in the array component of the result of+-- the lambda of a 'flatmap', which returns that array paired with a value of+-- regular type.+flatLambdaIrregulars :: E.Exp -> Int+flatLambdaIrregulars (E.Parens e _) = flatLambdaIrregulars e+flatLambdaIrregulars (E.Lambda _ _ _ (Info (E.RetType _ t)) _)+  | Just (irreg_t : _) <- E.isTupleRecord t =+      internalisedTypeSize $ E.toStruct irreg_t+flatLambdaIrregulars e =+  error $ "flatLambdaIrregulars: unexpected expression:\n" ++ prettyString e++-- | Internalise the lambda of a 'flatmap'. The irregular results (the ones that+-- are concatenated) are distinguished in the 'ExtLambda' by having the+-- existential size as their outermost size, so that size must be a variable -+-- when the lambda produces arrays of some other size, we bind that size and+-- coerce the arrays to it. By the source language type rules, this cannot fail.+internaliseFlatLambda :: E.Exp -> [Type] -> InternaliseM (I.ExtLambda SOACS)+internaliseFlatLambda lam argtypes = do+  (params, body, _) <- internaliseLambda lam argtypes+  (body', ret) <- buildBody . localScope (scopeOfLParams params) $ do+    res <- bodyBind body+    let (irreg_res, reg_res) = splitAt (flatLambdaIrregulars lam) res+    k <- irregularSize irreg_res+    irreg_res' <- mapM (coerceOuter k) irreg_res+    irreg_ts <- mapM (subExpType . resSubExp) irreg_res'+    reg_ts <- mapM (subExpType . resSubExp) reg_res+    pure+      ( subExpRes (I.Var k) : irreg_res' <> reg_res,+        I.Prim int64+          : existentialiseExtTypes [k] (staticShapes irreg_ts)+            <> staticShapes reg_ts+      )+  pure $ I.Lambda params ret body'+  where+    irregularSize [] =+      error "internaliseFlatLambda: lambda has no irregular results."+    irregularSize (r : _) =+      asVar "flatmap_k" . arraySize 0 =<< subExpType (resSubExp r)++    coerceOuter k r@(SubExpRes cs (I.Var v)) = do+      t <- lookupType v+      case arrayDims t of+        d : ds+          | d /= I.Var k ->+              SubExpRes cs+                <$> letSubExp "flatmap_irreg" (shapeCoerce (I.Var k : ds) v)+        _ -> pure r+    coerceOuter _ r = pure r+ -- | Overloaded operators are treated here. isOverloadedFunction ::   E.QualName VName ->@@ -1856,6 +1931,12 @@       internaliseHist 2 desc rf dest op ne buckets img     handleSOACs [rf, dest, op, ne, buckets, img] "hist_3d" = Just $ \desc ->       internaliseHist 3 desc rf dest op ne buckets img+    handleSOACs [lam, arr] "flatmap" = Just $ \desc -> do+      arr' <- internaliseExpToVars "map_arr" arr+      arr_ts <- mapM lookupType arr'+      lam' <- internaliseFlatLambda lam $ map rowType arr_ts+      let w = arraysSize 0 arr_ts+      letValExp' desc $ I.Op $ FlatMap w arr' lam'     handleSOACs _ _ = Nothing      handleAccs [dest, f, bs] "scatter_stream" = Just $ \desc ->@@ -2139,7 +2220,7 @@    shapeargs <- argShapes shapes fun_params argts   let diets =-        replicate (length shapeargs) I.ObservePrim+        replicate (length shapeargs) I.Observe           ++ map I.diet value_paramts   args' <-     ensureArgShapes@@ -2179,7 +2260,7 @@   ses_ts <- mapM subExpType ses    let combine t1 t2 =-        mconcat $ zipWith combine' (arrayExtDims t1) (arrayDims t2)+        mconcat $ zipWith combine' (arrayDims t1) (arrayDims t2)       combine' (I.Free (I.Var v)) se         | v `elem` retext = M.singleton v se       combine' _ _ = mempty
src/Futhark/Internalise/Monad.hs view
@@ -13,8 +13,8 @@     addFunDef,     lookupFunction,     lookupConst,-    bindFunction,-    bindConstant,+    addFunInfo,+    addConstant,     assert,     locating, @@ -171,13 +171,15 @@     (True, _) -> pure $ Just [Var fname]     _ -> pure Nothing -bindFunction :: VName -> FunDef SOACS -> FunInfo -> InternaliseM ()-bindFunction fname fd info = do-  addFunDef fd+-- | Register the calling information for a function, but not its+-- definition.  This is used to make a function available for+-- (recursive) calls before its body has been internalised.+addFunInfo :: VName -> FunInfo -> InternaliseM ()+addFunInfo fname info =   modify $ \s -> s {stateFunTable = M.insert fname info $ stateFunTable s} -bindConstant :: VName -> FunDef SOACS -> InternaliseM ()-bindConstant cname fd = do+addConstant :: VName -> FunDef SOACS -> InternaliseM ()+addConstant cname fd = do   addStms $ bodyStms $ funDefBody fd    case map resSubExp . bodyResult . funDefBody $ fd of
src/Futhark/Internalise/Monomorphise.hs view
@@ -178,6 +178,7 @@ data MonoState = MonoState   { sVNameSource :: !VNameSource,     sExpReplacements :: ExpReplacements,+    sAnySizes :: [(ReplacedExp, Int)],     sLifts :: Lifts,     sLiftedNames :: S.Set VName,     sValBinds :: [ValBind]@@ -205,7 +206,7 @@   )   where     ((), final_state) = runState (runReaderT m initial_env) initial_state-    initial_state = MonoState src mempty mempty mempty mempty+    initial_state = MonoState src mempty mempty mempty mempty mempty     initial_env = Env mempty mempty mempty mempty  lookupFun :: VName -> MonoM (Maybe PolyBinding)@@ -361,6 +362,16 @@ sizeVarName :: Exp -> Name sizeVarName e = "d<{" <> nameFromText (prettyText (bareExp e)) <> "}>" +-- | Is this size expression already normalised (a single variable or+-- constant), modulo being surrounded by some parentheses? If so, return the+-- normalised form.+maybeNormalisedSize :: Exp -> Maybe Exp+maybeNormalisedSize e+  | Just e' <- stripExp e = maybeNormalisedSize e'+maybeNormalisedSize (Var qn _ loc) = Just $ sizeFromName qn loc+maybeNormalisedSize (IntLit v _ loc) = Just $ IntLit v (Info i64) loc+maybeNormalisedSize _ = Nothing+ -- | Creates a new expression replacement if needed, this always produces normalised sizes. -- (e.g. single variable or constant) replaceExp :: Exp -> MonoM Exp@@ -378,17 +389,38 @@           vn <- newVName $ sizeVarName e           putExpReplacements . ((e', vn) :) =<< getExpReplacements           pure $ sizeFromName (qualName vn) (srclocOf e)-  where-    -- Avoid replacing of some 'already normalised' sizes that are just surounded by some parentheses.-    maybeNormalisedSize e'-      | Just e'' <- stripExp e' = maybeNormalisedSize e''-    maybeNormalisedSize (Var qn _ loc) = Just $ sizeFromName qn loc-    maybeNormalisedSize (IntLit v _ loc) = Just $ IntLit v (Info i64) loc-    maybeNormalisedSize _ = Nothing +-- | Turn a non-trivial size expression occurring in a *type* into an+-- 'anySize', rather than trying to hoist it into a named size parameter that+-- 'inferSizeArgs' would later have to (fallibly) reconstruct a concrete+-- argument for at every reference to the function. Some such expressions+-- (e.g. a product of several of the function's own parameters, as in #2230)+-- cannot in general be recovered at the point of reference, so instead of+-- pretending we can, we erase the information here and let internalisation+-- resolve it from the actual sizes of the values involved, exactly as is+-- already done for other 'anySize's (see Note [AnySize]).+--+-- Repeated occurrences of the same expression are mapped to the same+-- equivalence class, so that size equalities implied by the original+-- expression are not lost.+anySizeForType :: Exp -> MonoM Exp+anySizeForType e = do+  let e' = ReplacedExp e+  substs <- gets sAnySizes+  case lookup e' substs of+    Just i -> pure $ anySize i+    Nothing -> do+      i <- baseTag <$> newVName "any_size"+      modify $ \s -> s {sAnySizes = (e', i) : substs}+      pure $ anySize i+ transformFName :: SrcLoc -> QualName VName -> StructType -> MonoM Exp transformFName loc fname ft = do-  t' <- transformType ft+  -- The type 't' is later matched structurally, syntactically, against the+  -- callee's own declared signature (in 'inferSizeArgs', via 'dimMapping'),+  -- so it must stay fully precise here, even if we are currently inside the+  -- return type of some outer higher-order parameter.+  t' <- transformFNameType ft   let mono_t = monoType ft   if isIntrinsic (qualLeaf fname)     then pure $ var fname t'@@ -403,9 +435,7 @@         (Nothing, Nothing) -> pure $ var fname t'         -- A polymorphic function.         (Nothing, Just funbind) -> do-          (fname', infer, funbind') <- monomorphiseBinding funbind mono_t-          addValBind funbind'-          addLifted (qualLeaf fname) mono_t (fname', infer)+          (fname', infer) <- monomorphiseBinding funbind mono_t           applySizeArgs fname' (toRes Nonunique t') <$> infer t'   where     var fname' t' = Var fname' (Info t') loc@@ -430,8 +460,54 @@           )           size_args +-- | General-purpose size transformation for a type, used everywhere except when+-- reconstructing the type of a function reference (see 'transformFNameType').+-- Sizes immediately in the type (before any 'Arrow' is crossed) are hoisted+-- into a reference to a fresh named size variable (see 'replaceExp'), while+-- sizes found within a higher-order (function-valued) parameter or return+-- type - i.e. anywhere 'traverseDims' reports a 'DimPos' other than+-- 'PosImmediate' - are instead erased into an 'anySize'. Whether such a+-- hoisted name ends up as an actual explicit parameter of the enclosing+-- function, an ordinary local @let@-binding, or something else again,+-- depends on what the caller of 'transformType' does with the resulting+-- 'ExpReplacements' afterwards; 'transformType' itself does not add any+-- parameters. See Note [Higher-Order Parameter Sizes]. transformType :: TypeBase Size u -> MonoM (TypeBase Size u)-transformType typ =+transformType = traverseDims onDim+  where+    onDim _ pos e+      | Just _ <- isAnySize e = pure e+      | otherwise = do+          e' <- transformExp e+          case maybeNormalisedSize e' of+            Just e'' -> pure e''+            Nothing+              | PosImmediate <- pos -> replaceExp e'+              | otherwise -> anySizeForType e'++-- | Like 'transformType', but replaces every non-trivial size, regardless of+-- where it occurs, with a reference to a fresh named size variable (see+-- 'replaceExp'), rather than erasing sizes found within higher-order+-- parameter/return types into an 'anySize'. Note that this does not turn any+-- of these sizes into an actual parameter of anything: the resulting type has+-- exactly the same 'Arrow' structure (and hence the same arity) as the input.+-- The named variables are merely placeholders, used only so that+-- 'dimMapping' can structurally line this reconstructed type up against the+-- callee's own declared signature - whose *own* analogous sizes did, when+-- the callee itself was monomorphised, become real explicit parameters - so+-- that 'inferSizeArgs' can read off a concrete argument expression for each+-- of those parameters. Used only to reconstruct the type of a plain function+-- *reference* (in 'transformFName'); that is why it must stay fully precise,+-- even through further nested 'Arrow's - which, for this particular type,+-- are simply the callee's own remaining curried parameters, not necessarily+-- higher-order values. This also means we cannot use 'traverseDims' here+-- (unlike in 'transformType'): a size hoisted from deep within this type may+-- refer to a curried parameter bound by some enclosing 'Arrow' in the very+-- same type, so each 'Arrow' we cross must, before returning, re-scope any+-- such size locally to its own return type (via 'transformRetTypeSizesWith').+-- See Note [Higher-Order Parameter Sizes].+transformFNameType :: TypeBase Size u -> MonoM (TypeBase Size u)+transformFNameType typ =   case typ of     Scalar scalar ->       Scalar <$> transformScalarSizes scalar@@ -440,13 +516,13 @@   where     transformScalarSizes :: ScalarTypeBase Size u -> MonoM (ScalarTypeBase Size u)     transformScalarSizes (Record fs) =-      Record <$> traverse transformType fs+      Record <$> traverse transformFNameType fs     transformScalarSizes (Sum cs) =-      Sum <$> (traverse . traverse) transformType cs+      Sum <$> (traverse . traverse) transformFNameType cs     transformScalarSizes (Arrow as argName d argT retT) =       Arrow as argName d-        <$> transformType argT-        <*> transformRetTypeSizes argset retT+        <$> transformFNameType argT+        <*> transformRetTypeSizesWith transformFNameType argset retT       where         argset =           case argName of@@ -456,20 +532,33 @@       TypeVar u qn <$> mapM onArg args       where         onArg (TypeArgDim dim) = TypeArgDim <$> onDim dim-        onArg (TypeArgType ty) = TypeArgType <$> transformType ty+        onArg (TypeArgType ty) = TypeArgType <$> transformFNameType ty     transformScalarSizes ty@Prim {} = pure ty      onDim e       | Just _ <- isAnySize e = pure e       | otherwise = replaceExp =<< transformExp e -transformRetTypeSizes :: S.Set VName -> RetTypeBase Size as -> MonoM (RetTypeBase Size as)-transformRetTypeSizes argset (RetType dims ty) = do-  ty' <- withArgs argset $ withMono dims $ transformType ty+-- | Transform the sizes of a return type using the given size transformer,+-- given the set of parameter names newly in scope for this return type (usually+-- the name of the 'Arrow' parameter it is the return type of). Any size hoisted+-- while transforming that turns out to depend on one of these parameters cannot+-- be a top-level named parameter (those names are not in scope outside), so it+-- is instead added to the return type's own existentially-bound sizes.+transformRetTypeSizesWith ::+  (TypeBase Size as -> MonoM (TypeBase Size as)) ->+  S.Set VName ->+  RetTypeBase Size as ->+  MonoM (RetTypeBase Size as)+transformRetTypeSizesWith f argset (RetType dims ty) = do+  ty' <- withArgs argset $ withMono dims $ f ty   rl <- parametrizing argset   let dims' = dims <> map snd rl   pure $ RetType dims' ty' +transformRetTypeSizes :: S.Set VName -> RetTypeBase Size as -> MonoM (RetTypeBase Size as)+transformRetTypeSizes = transformRetTypeSizesWith transformType+ sizesForPat :: (MonadFreshNames m) => Pat ParamType -> m ([VName], Pat ParamType) sizesForPat pat = do   (params', sizes) <- runStateT (traverse (bitraverse onDim pure) pat) []@@ -1001,13 +1090,14 @@ removeEntryPoint (PolyBinding (_, name, tparams, params, rettype, body, attrs, loc)) =   PolyBinding (Nothing, name, tparams, params, rettype, body, attrs, loc) --- Monomorphise a polymorphic function at the types given in the instance--- list. Monomorphises the body of the function as well. Returns the fresh name--- of the generated monomorphic function and its 'ValBind' representation.+-- Monomorphise a polymorphic function at the types given in the instance list.+-- Monomorphises the body of the function as well. Returns the fresh name of the+-- generated monomorphic function as well a function for constructing additional+-- size arguments. monomorphiseBinding ::   PolyBinding ->   MonoType ->-  MonoM (VName, InferSizeArgs, ValBind)+  MonoM (VName, InferSizeArgs) monomorphiseBinding (PolyBinding (entry, name, tparams, params, rettype, body, attrs, loc)) inst_t = isolateNormalisation $ do   let bind_t = funType params rettype   (substs, t_shape_params) <-@@ -1042,41 +1132,49 @@        bind_t'' = funType params'' rettype''       bind_r = exp_naming <> extNaming-  body' <- updateExpTypes (`M.lookup` substs') body-  body'' <- withParams exp_naming' $ withArgs (shape_names <> args) $ transformExp body'-  scope' <- S.union (shape_names <> args) <$> askScope'-  body''' <--    expReplace exp_naming' <$> (calculateDims body'' . canCalculate scope' =<< getExpReplacements) +  -- It is important that we record the lifted function before transforming the+  -- body, in order to handle recursion. Fortunately the "calling convention"+  -- does not depend on the body.   seen_before <- elem name . map fst . M.keys <$> getLifts   name' <-     if null tparams && isNothing entry && not seen_before       then pure name       else newName name -  pure-    ( name',-      -- If the function is an entry point, then it cannot possibly-      -- need any explicit size arguments (checked by type checker).-      if isJust entry-        then const $ pure []-        else inferSizeArgs shape_params_explicit bind_t'' bind_r,-      if isJust entry-        then-          toValBinding-            name'-            (shape_params_explicit ++ shape_params_implicit)-            params''-            rettype''-            (entryAssert exp_naming body''')-        else-          toValBinding-            name'-            shape_params_implicit-            (map shapeParam shape_params_explicit ++ params'')-            rettype''-            body'''-    )+  let infer =+        -- If the function is an entry point, then it cannot possibly+        -- need any explicit size arguments (checked by type checker).+        if isJust entry+          then const $ pure []+          else inferSizeArgs shape_params_explicit bind_t'' bind_r++  addLifted name inst_t (name', infer)++  body' <- updateExpTypes (`M.lookup` substs') body+  body'' <- withParams exp_naming' $ withArgs (shape_names <> args) $ transformExp body'+  scope' <- S.union (shape_names <> args) <$> askScope'+  body''' <-+    expReplace exp_naming' <$> (calculateDims body'' . canCalculate scope' =<< getExpReplacements)++  addValBind $+    if isJust entry+      then+        toValBinding+          name'+          (shape_params_explicit ++ shape_params_implicit)+          params''+          rettype''+          (entryAssert exp_naming body''')+      else+        toValBinding+          name'+          shape_params_implicit+          (map shapeParam shape_params_explicit ++ params'')+          rettype''+          body'''++  pure (name', infer)   where     askScope' = S.filter (`notElem` retDims rettype) <$> askScope @@ -1208,9 +1306,7 @@           funType (valBindParams valbind) $             unInfo $               valBindRetType valbind-    (name, infer, valbind'') <- monomorphiseBinding valbind' $ monoType t-    addValBind valbind''-    addLifted (valBindName valbind) (monoType t) (name, infer)+    void $ monomorphiseBinding valbind' $ monoType t    let global =         if null (valBindParams valbind)@@ -1263,3 +1359,84 @@               map fst $                 M.toList b     )++-- Note [Higher-Order Parameter Sizes]+--+-- Normally, when 'transformType' (via 'onDim') encounters a size that is some+-- non-trivial expression rather than a single variable or a constant (say, 'm *+-- n'), it hoists that expression into a fresh named size parameter, tracked via+-- 'ExpReplacements' ('replaceExp'). Any function that ends up with such a+-- parameter in its (monomorphised) signature has it turned into an explicit+-- leading argument, and every reference to that function must then be augmented+-- with a concrete argument for it. This augmentation happens in+-- 'inferSizeArgs', which tries to recover a concrete expression for the+-- parameter by structurally comparing the callee's own declared type against+-- the type at the point of reference ('dimMapping').+--+-- This works fine as long as the expression's free variables are all reliably+-- available, spelled out the same way, at every place the function is+-- referenced. But consider (#2230):+--+--   def f (n: i64) (m: i64) (g: f64 -> [m ** n]f64) = g 0+--   entry main n m = f n m (\x -> replicate (m ** n) x)+--+-- Here 'm ** n' only occurs as the return size of a *parameter* of 'f' (the+-- higher-order function 'g'), not as a size of 'f' itself. Because+-- 'transformFName' processes the bare reference to 'f' before its arguments+-- have been applied, 'dimMapping' has no way to know, at that point, that the+-- surrounding application happens to supply concrete values for 'n' and 'm'.+-- (In general it may not even be a full application - 'f' could be passed+-- onward as a value.) 'inferSizeArgs' then has no honest answer and used to+-- fall back to inserting a literal '0' as the argument, which will then be+-- treated as ground truth by the internaliser.+--+-- The same thing happens if the problematic size instead occurs in the *+-- argument* type of a returned or parameter-bound function, e.g.+--+--   def f (n: i64) (m: i64) =+--     \(g: [m ** n]f64 -> f64) -> g (replicate (m ** n) 0.0)+--   entry main n m = f n m (\x -> head x)+--+-- (see @test.fut@ in the repository root at the time of writing). 'd = m ** n'+-- is here unwitnessed by anything else in 'f's own signature, so it again+-- becomes an explicit parameter of 'f' that 'inferSizeArgs' cannot recover at+-- the unsaturated reference to 'f'.+--+-- The fix implemented here recognises that a size occurring anywhere in the+-- type of a higher-order (function-valued) parameter or return value - whether+-- describing what it consumes or what it produces - is different in kind from+-- an ordinary size: the function that declares it (here 'f') always has the+-- free variables it depends on in scope, since such a size can only be built+-- from parameters of the enclosing function itself. So nothing ever needs its+-- value threaded in externally as an explicit argument; the enclosing function+-- can always reconstruct it itself, wherever it actually needs it (to type the+-- result of calling 'g', to build an array to pass to 'g', or just to pass it+-- on to a further nested 'Arrow'). We erase the expression into a fresh+-- 'anySize' right where it occurs, instead of pretending we can reconstruct it+-- as an explicit argument (see Note [AnySize]). Since the size is never turned+-- into an explicit parameter of 'f' in the first place, there is also nothing+-- for 'inferSizeArgs' to fail to supply.+--+-- Deciding when we are in "the type of a higher-order parameter or return+-- value" is the job of the 'DimPos' that 'traverseDims' reports for each+-- size it visits: 'PosImmediate' means the size is not (yet) inside any+-- 'Arrow', while 'PosParam'/'PosReturn' mean it is on the argument or return+-- side of one (we treat the two identically). 'transformType' is thus+-- simply 'traverseDims' with a callback that hoists at 'PosImmediate' and+-- erases to an 'anySize' otherwise; unlike the hand-rolled traversal this+-- replaced, it gets the structural recursion through 'Record'/'Sum'/'Arrow'+-- for free. The one place that cannot use 'traverseDims' directly is+-- 'transformFNameType', which must stay precise through every nested+-- 'Arrow' regardless of position (see its docstring for why).+--+-- Repeated occurrences of the same original expression within one binding group+-- (e.g. the same higher-order parameter's size mentioned twice, or two+-- different higher-order parameters sharing a size) are mapped to the same+-- 'anySize' equivalence class via the 'sAnySizes' cache in 'MonoState', rather+-- than each minting a fresh one, mirroring 'sizesForPat' below. Two erasures+-- that end up in *different* equivalence classes (e.g. one introduced while+-- processing 'f's own parameter list, another introduced independently at some+-- call site's reconstructed argument type) are not a problem: 'anySize'+-- equivalence classes are only ever compared within a single binding group,+-- never across independent functions or call sites, so there is nothing to keep+-- synchronised between the two.
src/Futhark/Internalise/TypesValues.hs view
@@ -64,17 +64,22 @@     onType = fromMaybe bad . hasStaticShape     bad = error $ "internaliseParamTypes: " ++ prettyString ts --- We need to fix up the arrays for any Acc return values or loop--- parameters.  We look at the concrete types for this, since the Acc--- parameter name in the second list will just be something we made up.+-- Replace an accumulator's token, index space, and element types with those of+-- a known accumulator type. We must do this because these components cannot be+-- recovered from a source type: 'internaliseTypeM' produces a placeholder token+-- and a guessed index space. The known type is computed elsewhere (from+-- concrete loop values, or from an accumulator parameter).+fixupAcc :: TypeBase shape1 u1 -> (TypeBase shape2 u2, b) -> (TypeBase shape2 u2, b)+fixupAcc (Acc acc ispace ts _) (Acc _ _ _ u, b) = (Acc acc ispace ts u, b)+fixupAcc _ t = t++-- Fix up accumulators using a positionally-matching list of concrete+-- types (e.g. the actual types of loop values). fixupKnownTypes ::   [TypeBase shape1 u1] ->   [(TypeBase shape2 u2, b)] ->   [(TypeBase shape2 u2, b)]-fixupKnownTypes = zipWith fixup-  where-    fixup (Acc acc ispace ts _) (Acc _ _ _ u2, b) = (Acc acc ispace ts u2, b)-    fixup _ t = t+fixupKnownTypes = zipWith fixupAcc  -- Generate proper certificates for the placeholder accumulator -- certificates produced by internaliseType (identified with tag 0).@@ -163,13 +168,17 @@ internaliseReturnType ::   [Tree (I.TypeBase Shape Uniqueness)] ->   E.ResRetType ->-  [TypeBase shape u] ->   [(I.TypeBase ExtShape Uniqueness, RetAls)]-internaliseReturnType paramts (E.RetType dims et) ts =-  fixupKnownTypes ts . concat . inferAliases paramts $+internaliseReturnType paramts (E.RetType dims et) =+  fixupAccs . concat . inferAliases paramts $     runInternaliseTypeM' dims (internaliseTypeM exts et)   where     exts = M.fromList $ zip dims [0 ..]+    -- Any 'Acc' in the return type must (by the type rules) be the function's+    -- single 'Acc' parameter, so we substitute its known accumulator type.+    fixupAccs = case [t | t@Acc {} <- foldMap toList paramts] of+      acc : _ -> map (fixupAcc acc)+      [] -> id  -- | As 'internaliseReturnType', but returns components of a top-level -- tuple type piecemeal.@@ -188,10 +197,9 @@  internaliseCoerceType ::   E.StructType ->-  [TypeBase shape u] ->   [I.TypeBase ExtShape Uniqueness]-internaliseCoerceType et ts =-  map fst $ internaliseReturnType [] (E.RetType [] $ E.toRes E.Nonunique et) ts+internaliseCoerceType et =+  map fst $ internaliseReturnType [] (E.RetType [] $ E.toRes E.Nonunique et)  internaliseLambdaReturnType ::   E.ResType ->
src/Futhark/Optimise/ArrayShortCircuiting.hs view
@@ -122,9 +122,12 @@   elems' <- mapM replaceInPatElem elems   e' <- replaceInExp elems' e   entries <- asks (M.elems . envCoalesceTab)-  let c' = case filter (\entry -> (map patElemName elems `L.intersect` M.keys (vartab entry)) /= []) entries of+  let bound_here = map patElemName elems+      -- Also remove certs produced here to avoid circularity.+      inScopeCerts = Certs . filter (`notElem` bound_here) . unCerts+      c' = case filter (\entry -> (bound_here `L.intersect` M.keys (vartab entry)) /= []) entries of         [] -> c-        entries' -> c <> foldMap certs entries'+        entries' -> c <> inScopeCerts (foldMap certs entries')   pure $ Let (Pat elems') (StmAux c' a loc d) e'   where     replaceInPatElem :: PatElem LetDecMem -> UpdateM inner (PatElem LetDecMem)
src/Futhark/Optimise/BlkRegTiling.hs view
@@ -274,7 +274,7 @@     isAcc res_tp,     -- we get the global-thread id for the two inner dimensions,     --   as we are probably going to use it in code generation-    (gtid_x, width_B) : (gtid_y, height_A) : rem_outer_dims_rev <-+    (gtid_y, width_B) : (gtid_x, height_A) : rem_outer_dims_rev <-       reverse $ unSegSpace seg_space,     rem_outer_dims <- reverse rem_outer_dims_rev,     Just@@ -298,10 +298,10 @@         tk_rk <- letSubExp "tk_rk" =<< toExp (pe64 tk * pe64 rk)          gridDim_t <- letSubExp "gridDim_t" =<< ceilDiv common_dim tk_rk-        gridDim_y <- letSubExp "gridDim_y" =<< ceilDiv height_A ty_ry-        gridDim_x <- letSubExp "gridDim_x" =<< ceilDiv width_B tx_rx+        gridDim_x <- letSubExp "gridDim_x" =<< ceilDiv height_A tx_rx+        gridDim_y <- letSubExp "gridDim_y" =<< ceilDiv width_B ty_ry -        let gridxyt_pexp = pe64 gridDim_y * pe64 gridDim_x * pe64 gridDim_t+        let gridxyt_pexp = pe64 gridDim_x * pe64 gridDim_y * pe64 gridDim_t             grid_pexp =               foldl (\x d -> pe64 d * x) gridxyt_pexp $                 map snd rem_outer_dims_rev@@ -312,8 +312,8 @@          ---- in this binder: outer seggroup ----         (ret_seggroup, stms_seggroup) <- runBuilder $ do-          iii <- letExp "iii" =<< toExp (le64 gid_y * pe64 ty_ry)-          jjj <- letExp "jjj" =<< toExp (le64 gid_x * pe64 tx_rx)+          iii <- letExp "iii" =<< toExp (le64 gid_x * pe64 tx_rx)+          jjj <- letExp "jjj" =<< toExp (le64 gid_y * pe64 ty_ry)           ttt <- letExp "ttt" =<< toExp (le64 gid_t * pe64 tk_rk)            -- initialize register mem with neutral elements and create shmem@@ -333,7 +333,7 @@                 ( (rx, ry, tx, ty, tk, tk_div_tx, tk_div_ty, tx_rx),                   segthd_lvl,                   var_dims,-                  (gtid_x, width_B, gtid_y, height_A, common_dim),+                  (gtid_y, width_B, gtid_x, height_A, common_dim),                   (iii, jjj),                   (load_A, inp_A, map_t1, load_B, inp_B, map_t2),                   (map_lam, red_lam)@@ -382,13 +382,13 @@             (res_nm, res_tp)             (ty, tx, ry, rx)             (iii, jjj)-            (gtid_y, gtid_x)+            (gtid_x, gtid_y)             (height_A, width_B, rem_outer_dims)             code2'          let grid = KernelGrid (Count grid_size) (Count tblock_size)             level' = SegBlock SegNoVirt (Just grid)-            space' = SegSpace gid_flat (rem_outer_dims ++ [(gid_t, gridDim_t), (gid_y, gridDim_y), (gid_x, gridDim_x)])+            space' = SegSpace gid_flat (rem_outer_dims ++ [(gid_t, gridDim_t), (gid_x, gridDim_x), (gid_y, gridDim_y)])             kbody' = Body () stms_seggroup ret_seggroup         pure $ Let pat aux $ Op $ SegOp $ SegMap level' space' ts kbody'       pure $ Just (host_stms, new_kernel)@@ -469,7 +469,7 @@     primType res_tp,     -- we get the global-thread id for the two inner dimensions,     --   as we are probably going to use it in code generation-    (gtid_x, width_B) : (gtid_y, height_A) : rem_outer_dims_rev <-+    (gtid_y, width_B) : (gtid_x, height_A) : rem_outer_dims_rev <-       reverse $ unSegSpace seg_space,     rem_outer_dims <- reverse rem_outer_dims_rev,     Just@@ -488,9 +488,9 @@         (rx, ry, tx, ty, tk, tk_div_tx, tk_div_ty, tx_rx, ty_ry, a_loc_sz, b_loc_sz) <-           mkTileMemSizes height_A width_B common_dim is_B_coal -        gridDim_x <- letSubExp "gridDim_x" =<< ceilDiv width_B tx_rx-        gridDim_y <- letSubExp "gridDim_y" =<< ceilDiv height_A ty_ry-        let gridxy_pexp = pe64 gridDim_y * pe64 gridDim_x+        gridDim_y <- letSubExp "gridDim_y" =<< ceilDiv width_B ty_ry+        gridDim_x <- letSubExp "gridDim_x" =<< ceilDiv height_A tx_rx+        let gridxy_pexp = pe64 gridDim_x * pe64 gridDim_y         let grid_pexp =               foldl (\x d -> pe64 d * x) gridxy_pexp $                 map snd rem_outer_dims_rev@@ -500,8 +500,8 @@          ---- in this binder: outer seggroup ----         (ret_seggroup, stms_seggroup) <- runBuilder $ do-          iii <- letExp "iii" =<< toExp (le64 gid_y * pe64 ty_ry)-          jjj <- letExp "jjj" =<< toExp (le64 gid_x * pe64 tx_rx)+          iii <- letExp "iii" =<< toExp (le64 gid_x * pe64 tx_rx)+          jjj <- letExp "jjj" =<< toExp (le64 gid_y * pe64 ty_ry)            -- initialize register mem with neutral elements and create shmem           (cssss, a_loc_init, b_loc_init) <-@@ -521,7 +521,7 @@                 ( (rx, ry, tx, ty, tk, tk_div_tx, tk_div_ty, tx_rx),                   segthd_lvl,                   var_dims,-                  (gtid_x, width_B, gtid_y, height_A, common_dim),+                  (gtid_y, width_B, gtid_x, height_A, common_dim),                   (iii, jjj),                   (load_A, inp_A, map_t1, load_B, inp_B, map_t2),                   (map_lam, red_lam)@@ -543,7 +543,7 @@           let redomap_res : _ = epilogue_res_list            -- support for non-empty code2'-          --  segmap (ltid_y < ty, ltid_x < tx) {+          --  segmap (ltid_x < tx, ltid_y < ty) {           --    for i < ry do           --      for j < rx do           --        res = if (iii+ltid_y*ry+i < height_A && jjj+ltid_x*rx+j < width_B)@@ -555,13 +555,13 @@             (res_nm, res_tp)             (ty, tx, ry, rx)             (iii, jjj)-            (gtid_y, gtid_x)+            (gtid_x, gtid_y)             (height_A, width_B, rem_outer_dims)             code2'          let grid = KernelGrid (Count grid_size) (Count tblock_size)             level' = SegBlock SegNoVirt (Just grid)-            space' = SegSpace gid_flat (rem_outer_dims ++ [(gid_y, gridDim_y), (gid_x, gridDim_x)])+            space' = SegSpace gid_flat (rem_outer_dims ++ [(gid_x, gridDim_x), (gid_y, gridDim_y)])             kbody' = Body () stms_seggroup ret_seggroup         pure $ Let pat aux $ Op $ SegOp $ SegMap level' space' ts kbody'       pure $ Just (host_stms, new_kernel)@@ -1005,7 +1005,7 @@     variance <- varianceInStms initial_variance kstms,     -- we get the global-thread id for the two inner dimensions,     --   as we are probably going to use it in code generation-    (gtid_x, d_Kx) : (gtid_y, d_Ky) : (gtid_z, d_M) : rem_outer_dims_rev <- reverse $ unSegSpace space,+    (gtid_z, d_Kz) : (gtid_y, d_Ky) : (gtid_x, d_Kx) : rem_outer_dims_rev <- reverse $ unSegSpace space,     rem_outer_dims <- reverse rem_outer_dims_rev,     -- check that the code fits the pattern having:     -- some `code1`, followed by one Screma SOAC, followed by some `code2`@@ -1046,42 +1046,42 @@     length tmp_stms == length inp_soac_arrs,     -- code1' <- stmsFromList $ stmsToList code1 \\ stmsToList code2'',     code2' <- code2'' <> code2,-    -- we assume the kernel results are variant to the thrid-outer parallel dimension-    -- (for sanity sake, they should be)+    -- we assume the kernel results are variant to the outermost of the three+    -- innermost parallel dimensions (for sanity sake, they should be)     ker_res_nms <- mapMaybe getResNm kres,     length ker_res_nms == length kres,     all primType kertp,-    all (variantToDim variance gtid_z) ker_res_nms = do+    all (variantToDim variance gtid_x) ker_res_nms = do       -- HERE STARTS THE IMPLEMENTATION:       (new_kernel, host_stms) <- runBuilder $ do         -- host code-        -- process the z-variant arrays that need transposition;+        -- process the x-variant arrays that need transposition;         -- these "manifest" statements should come before the kernel         (tab_inn, tab_out) <-           foldM-            (insertTranspose variance (gtid_z, d_M))+            (insertTranspose variance (gtid_x, d_Kx))             (M.empty, M.empty)             $ M.toList arr_tab0 -        tx_name <- nameFromText . prettyText <$> newVName "Tx"+        tz_name <- nameFromText . prettyText <$> newVName "Tz"         ty_name <- nameFromText . prettyText <$> newVName "Ty" -        tx0 <- letSubExp "Tx" $ Op $ SizeOp $ GetSize tx_name SizeTile+        tz0 <- letSubExp "Tz" $ Op $ SizeOp $ GetSize tz_name SizeTile         ty0 <- letSubExp "Ty" $ Op $ SizeOp $ GetSize ty_name SizeTile+        tz <- limitTile "Tz" tz0 d_Kz         ty <- limitTile "Ty" ty0 d_Ky-        tx <- limitTile "Tx" tx0 d_Kx-        let rz = reg_tile_se+        let rx = reg_tile_se -        gridDim_x <- letSubExp "gridDim_x" =<< ceilDiv d_Kx tx+        gridDim_z <- letSubExp "gridDim_z" =<< ceilDiv d_Kz tz         gridDim_y <- letSubExp "gridDim_y" =<< ceilDiv d_Ky ty-        gridDim_z <- letSubExp "gridDim_z" =<< ceilDiv d_M rz-        let gridxyz_pexp = pe64 gridDim_z * pe64 gridDim_y * pe64 gridDim_x+        gridDim_x <- letSubExp "gridDim_x" =<< ceilDiv d_Kx rx+        let gridxyz_pexp = pe64 gridDim_x * pe64 gridDim_y * pe64 gridDim_z         let grid_pexp = product $ gridxyz_pexp : map (pe64 . snd) rem_outer_dims_rev         grid_size <- letSubExp "grid_size_tile3d" =<< toExp grid_pexp-        tblock_size <- letSubExp "tblock_size_tile3d" =<< toExp (pe64 ty * pe64 tx)+        tblock_size <- letSubExp "tblock_size_tile3d" =<< toExp (pe64 ty * pe64 tz)         let segthd_lvl = SegThreadInBlock (SegNoVirtFull (SegSeqDims [])) -        count_shmem <- letSubExp "count_shmem" =<< ceilDiv rz tblock_size+        count_shmem <- letSubExp "count_shmem" =<< ceilDiv rx tblock_size          gid_x <- newVName "gid_x"         gid_y <- newVName "gid_y"@@ -1090,15 +1090,15 @@          ---- in this binder: outer seggroup ----         (ret_seggroup, stms_seggroup) <- runBuilder $ do-          ii <- letExp "ii" =<< toExp (le64 gid_z * pe64 rz)+          ii <- letExp "ii" =<< toExp (le64 gid_x * pe64 rx)           jj1 <- letExp "jj1" =<< toExp (le64 gid_y * pe64 ty)-          jj2 <- letExp "jj2" =<< toExp (le64 gid_x * pe64 tx)+          jj2 <- letExp "jj2" =<< toExp (le64 gid_z * pe64 tz)            -- initialize the register arrays corresponding to the result of redomap;-          reg_arr_nms <- segMap2D "res" segthd_lvl ResultPrivate (ty, tx) $ \_ ->+          reg_arr_nms <- segMap2D "res" segthd_lvl ResultPrivate (ty, tz) $ \_ ->             forM (zip red_nes red_res_tps) $ \(red_ne, red_t) -> do-              css_init <- scratch "res_init" (elemType red_t) [rz]-              css <- forLoop rz [css_init] $ \i [css_merge] -> do+              css_init <- scratch "res_init" (elemType red_t) [rx]+              css <- forLoop rx [css_init] $ \i [css_merge] -> do                 css' <- update "css" css_merge [i] red_ne                 resultBodyM [Var css']               pure $ varRes css@@ -1106,7 +1106,7 @@           -- scratch the shared-memory arrays corresponding to the arrays that are           --   input to the redomap and are invariant to the outermost parallel dimension.           loc_arr_nms <- forM (M.toList tab_out) $ \(nm, (ptp, _)) ->-            scratch (baseName nm <> "_loc") ptp [rz]+            scratch (baseName nm <> "_loc") ptp [rx]            prologue_res_list <-             forLoop' common_dim (reg_arr_nms ++ loc_arr_nms) $@@ -1126,12 +1126,12 @@                           body <- runBodyBuilder $ do                             offs <- letExp "offs" =<< toExp (pe64 tblock_size * le64 tt)                             loc_ind <- letExp "loc_ind" =<< toExp (le64 ltid + le64 offs)-                            letBindNames [gtid_z] =<< toExp (le64 ii + le64 loc_ind)-                            let glb_ind = gtid_z+                            letBindNames [gtid_x] =<< toExp (le64 ii + le64 loc_ind)+                            let glb_ind = gtid_x                             y_elm <-                               letSubExp "y_elem"                                 =<< eIf-                                  (toExp $ le64 glb_ind .<. pe64 d_M)+                                  (toExp $ le64 glb_ind .<. pe64 d_Kx)                                   ( do                                       addStm load_Y                                       res <- index "Y_elem" glb_Y_nm [q]@@ -1141,7 +1141,7 @@                             y_ind <-                               letSubExp "y_loc_ind"                                 =<< eIf-                                  (toExp $ le64 loc_ind .<. pe64 rz)+                                  (toExp $ le64 loc_ind .<. pe64 rx)                                   (toExp loc_ind >>= letTupExp' "loc_fi" >>= resultBodyM)                                   (eBody [pure $ BasicOp $ SubExp $ intConst Int64 (-1)])                             acc' <- letExp (baseName acc) $ BasicOp $ UpdateAcc Safe acc [y_ind] [y_elm]@@ -1155,15 +1155,15 @@                     resultBodyM $ map Var loc_arr_merge2_nms'                  redomap_res <--                  segMap2D "redomap_res" segthd_lvl ResultPrivate (ty, tx) $-                    \(ltid_y, ltid_x) -> do+                  segMap2D "redomap_res" segthd_lvl ResultPrivate (ty, tz) $+                    \(ltid_y, ltid_z) -> do                       letBindNames [gtid_y] =<< toExp (le64 jj1 + le64 ltid_y)-                      letBindNames [gtid_x] =<< toExp (le64 jj2 + le64 ltid_x)+                      letBindNames [gtid_z] =<< toExp (le64 jj2 + le64 ltid_z)                       reg_arr_merge_nms_slc <- forM reg_arr_merge_nms $ \reg_arr_nm ->-                        index "res_reg_slc" reg_arr_nm [ltid_y, ltid_x]+                        index "res_reg_slc" reg_arr_nm [ltid_y, ltid_z]                       fmap subExpsRes . letTupExp' "redomap_guarded"                         =<< eIf-                          (toExp $ le64 gtid_y .<. pe64 d_Ky .&&. le64 gtid_x .<. pe64 d_Kx)+                          (toExp $ le64 gtid_y .<. pe64 d_Ky .&&. le64 gtid_z .<. pe64 d_Kz)                           ( do                               inp_scals_invar_outer <-                                 forM (M.toList tab_inn) $ \(inp_arr_nm, load_stm) -> do@@ -1171,18 +1171,18 @@                                   index (baseName inp_arr_nm) inp_arr_nm [q]                               -- build the loop of count R whose body is semantically the redomap code                               reg_arr_merge_nms' <--                                forLoop' rz reg_arr_merge_nms_slc $ \i reg_arr_mm_nms -> do-                                  letBindNames [gtid_z] =<< toExp (le64 ii + le64 i)+                                forLoop' rx reg_arr_merge_nms_slc $ \i reg_arr_mm_nms -> do+                                  letBindNames [gtid_x] =<< toExp (le64 ii + le64 i)                                   resultBodyM                                     =<< letTupExp' "redomap_lam"                                     =<< eIf-                                      (toExp $ le64 gtid_z .<. pe64 d_M)+                                      (toExp $ le64 gtid_x .<. pe64 d_Kx)                                       ( do                                           -- read from shared memory                                           ys <- forM loc_arr_nms' $ \loc_arr_nm ->-                                            index "inp_reg_var2z" loc_arr_nm [i]+                                            index "inp_reg_var2x" loc_arr_nm [i]                                           cs <- forM reg_arr_mm_nms $ \reg_arr_nm ->-                                            index "res_reg_var2z" reg_arr_nm [i]+                                            index "res_reg_var2x" reg_arr_nm [i]                                           -- here we need to put in order the scalar inputs to map:                                           let tab_scals =                                                 M.fromList $@@ -1207,33 +1207,33 @@                 resultBodyM $ map Var $ redomap_res ++ loc_arr_nms'            -- support for non-empty code2'-          --  segmap (ltid_y < ty, ltid_x < tx) {-          --    for i < rz do-          --        res = if (ii+i < d_M && jj1+ltid_y < d_Ky && jj2 + ltid_x < d_Kx)+          --  segmap (ltid_y < ty, ltid_z < tz) {+          --    for i < rx do+          --        res = if (ii+i < d_Kx && jj1+ltid_y < d_Ky && jj2 + ltid_z < d_Kz)           --              then code2' else dummy           --        final_res[i] = res           let redomap_res = take (length red_nes) prologue_res_list           epilogue_res <-             if length redomap_orig_res == length ker_res_nms               && ker_res_nms == map patElemName redomap_orig_res-              then segMap3D "rssss" segthd_lvl ResultPrivate (se1, ty, tx) $ \(_ltid_z, ltid_y, ltid_x) ->+              then segMap3D "rssss" segthd_lvl ResultPrivate (se1, ty, tz) $ \(_ltid_x, ltid_y, ltid_z) ->                 forM (zip kertp redomap_res) $ \(res_tp, res) -> do-                  rss_init <- scratch "rss_init" (elemType res_tp) [rz, se1, se1]+                  rss_init <- scratch "rss_init" (elemType res_tp) [rx, se1, se1]                   fmap varRes $-                    forLoop rz [rss_init] $ \i [rss] -> do+                    forLoop rx [rss_init] $ \i [rss] -> do                       let slice = Slice [DimFix $ Var i, DimFix se0, DimFix se0]-                      thread_res <- index "thread_res" res [ltid_y, ltid_x, i]+                      thread_res <- index "thread_res" res [ltid_y, ltid_z, i]                       rss' <- letSubExp "rss" $ BasicOp $ Update Unsafe rss slice $ Var thread_res                       resultBodyM [rss']-              else segMap3D "rssss" segthd_lvl ResultPrivate (se1, ty, tx) $ \(_ltid_z, ltid_y, ltid_x) -> do+              else segMap3D "rssss" segthd_lvl ResultPrivate (se1, ty, tz) $ \(_ltid_x, ltid_y, ltid_z) -> do                 letBindNames [gtid_y] =<< toExp (le64 jj1 + le64 ltid_y)-                letBindNames [gtid_x] =<< toExp (le64 jj2 + le64 ltid_x)+                letBindNames [gtid_z] =<< toExp (le64 jj2 + le64 ltid_z)                 rss_init <- forM kertp $ \res_tp ->-                  scratch "rss_init" (elemType res_tp) [rz, se1, se1]-                rss <- forLoop' rz rss_init $ \i rss_merge -> do-                  letBindNames [gtid_z] =<< toExp (le64 ii + le64 i)+                  scratch "rss_init" (elemType res_tp) [rx, se1, se1]+                rss <- forLoop' rx rss_init $ \i rss_merge -> do+                  letBindNames [gtid_x] =<< toExp (le64 ii + le64 i)                   forM_ (zip redomap_orig_res redomap_res) $ \(o_res, n_res) -> do-                    c <- index "redomap_thd" n_res [ltid_y, ltid_x, i]+                    c <- index "redomap_thd" n_res [ltid_y, ltid_z, i]                     letBindNames [patElemName o_res] =<< toExp (le64 c)                     pure c                   res_els <-@@ -1242,10 +1242,10 @@                         ( toExp $                             le64 gtid_y                               .<. pe64 d_Ky+                              .&&. le64 gtid_z+                              .<. pe64 d_Kz                               .&&. le64 gtid_x                               .<. pe64 d_Kx-                              .&&. le64 gtid_z-                              .<. pe64 d_M                         )                         ( do                             addStms code2'@@ -1263,7 +1263,7 @@           ----------------------------------------------------------------           let regtile_ret_dims =                 map (\(_, sz) -> (sz, se1, se1)) rem_outer_dims-                  ++ [(d_M, se1, rz), (d_Ky, ty, se1), (d_Kx, tx, se1)]+                  ++ [(d_Kx, se1, rx), (d_Ky, ty, se1), (d_Kz, tz, se1)]            epilogue_res' <- forM epilogue_res $ \res ->             if null rem_outer_dims@@ -1281,7 +1281,7 @@         -- END (ret_seggroup, stms_seggroup) <- runBuilder $ do         let grid = KernelGrid (Count grid_size) (Count tblock_size)             level' = SegBlock SegNoVirt (Just grid)-            space' = SegSpace gid_flat (rem_outer_dims ++ [(gid_z, gridDim_z), (gid_y, gridDim_y), (gid_x, gridDim_x)])+            space' = SegSpace gid_flat (rem_outer_dims ++ [(gid_x, gridDim_x), (gid_y, gridDim_y), (gid_z, gridDim_z)])             kbody' = Body () stms_seggroup ret_seggroup          pure $ Let pat aux $ Op $ SegOp $ SegMap level' space' kertp kbody'
src/Futhark/Optimise/CSE.hs view
@@ -156,8 +156,8 @@  cseInLambda ::   (Aliased rep, CSEInOp (Op rep)) =>-  Lambda rep ->-  CSEM rep (Lambda rep)+  GLambda rep t ->+  CSEM rep (GLambda rep t) cseInLambda lam = do   body' <- cseInBody (map (const Observe) $ lambdaReturnType lam) $ lambdaBody lam   pure lam {lambdaBody = body'}@@ -331,4 +331,4 @@   (AliasableRep rep, CSEInOp (Op (Aliases rep))) =>   CSEInOp (SOAC.SOAC (Aliases rep))   where-  cseInOp = subCSE . SOAC.mapSOACM (SOAC.SOACMapper pure cseInLambda pure)+  cseInOp = subCSE . SOAC.mapSOACM (SOAC.SOACMapper pure cseInLambda cseInLambda pure)
src/Futhark/Optimise/DoubleBuffer.hs view
@@ -149,9 +149,12 @@           mapOnOp = onOp         } +optLoops :: (Constraints rep inner) => DoubleBufferM rep m -> DoubleBufferM rep m+optLoops = local $ \env -> env {envOptimiseLoop = optimiseLoop}+ optimiseGPUOp :: OptimiseOp GPUMem optimiseGPUOp (Inner (SegOp op)) =-  local inSegOp $ Inner . SegOp <$> mapSegOpM mapper op+  optLoops $ Inner . SegOp <$> mapSegOpM mapper op   where     mapper =       identitySegOpMapper@@ -159,12 +162,13 @@           mapOnSegPostOpLambda = optimiseLambda,           mapOnSegOpBody = optimiseKernelBody         }-    inSegOp env = env {envOptimiseLoop = optimiseLoop}+optimiseGPUOp (Inner (GPUBody ts body)) =+  optLoops $ Inner . GPUBody ts <$> optimiseBody body optimiseGPUOp op = pure op  optimiseMCOp :: OptimiseOp MCMem optimiseMCOp (Inner (ParOp par_op op)) =-  local inSegOp $+  optLoops $     Inner       <$> (ParOp <$> traverse (mapSegOpM mapper) par_op <*> mapSegOpM mapper op)   where@@ -174,7 +178,6 @@           mapOnSegPostOpLambda = optimiseLambda,           mapOnSegOpBody = optimiseKernelBody         }-    inSegOp env = env {envOptimiseLoop = optimiseLoop} optimiseMCOp op = pure op  optimiseKernelBody ::
src/Futhark/Optimise/Fusion/GraphRep.hs view
@@ -395,6 +395,7 @@   Futhark.JVP {} -> freeClassifications soac   Futhark.VJP {} -> freeClassifications soac   Futhark.WithVJP {} -> freeClassifications soac+  Futhark.FlatMap {} -> freeClassifications soac   where     inputs = S.fromList . (`zip` repeat SOACInput) expInputs e
src/Futhark/Optimise/Fusion/Screma.hs view
@@ -102,9 +102,6 @@   when     (parAccsOverlap new_lam new_lam')     (fail "Can not fuse due to overlap in parameter accumalators.")-  when-    (resAccsOverlap new_lam new_lam')-    (fail "Can not fuse due to overlap in result accumalators.")   pure     ( (new_inp, new_lam, new_out),       (new_inp', new_lam', new_out')@@ -134,16 +131,14 @@ -- consumers scans or reduces. fusible ::   (MonadFail m) =>-  [SOAC.Input] ->   ScremaForm SOACS ->   [VName] ->   [SOAC.Input] ->   ScremaForm SOACS ->-  [VName] ->   m ()-fusible inp_p form_p out_p inp_c form_c out_c = do+fusible form_p out_p inp_c form_c = do   ((_, post_scan_p, _), _) <--    splitLambdaByPar post_scan_pars_p inp_p post_p out_c+    splitLambdaByPar post_scan_pars_p (lambdaParams post_p) post_p (lambdaReturnType post_p)   let post_scan_res_p = bodyResult $ lambdaBody post_scan_p       forbidden_p = namesFromList $ resToOut out_p post_p <$> post_scan_res_p       is_fusible =@@ -381,22 +376,12 @@ parAccs :: Lambda SOACS -> [Type] parAccs = filter isAcc . map typeOf . lambdaParams --- | Find all Accumulator results.-resAccs :: Lambda SOACS -> [Type]-resAccs = filter isAcc . lambdaReturnType- -- | Check if the lambda parameters have overlapping accumulators. parAccsOverlap :: Lambda SOACS -> Lambda SOACS -> Bool parAccsOverlap lam = any (`elem` accs) . parAccs   where     accs = parAccs lam --- | Check if the lambdas result have overlapping accumulators.-resAccsOverlap :: Lambda SOACS -> Lambda SOACS -> Bool-resAccsOverlap lam = any (`elem` accs) . resAccs-  where-    accs = resAccs lam- -- | Check if the lambdas have parameters that overlap due to -- consumption. consumedOverlap :: Lambda SOACS -> Lambda SOACS -> Bool@@ -496,7 +481,7 @@   [VName] ->   m ([SOAC.Input], ScremaForm SOACS, [VName]) fuseScrema w inp_p form_p out_p inp_c form_c out_c = do-  fusible inp_p form_p out_p inp_c form_c out_c+  fusible form_p out_p inp_c form_c   (super_screma, new_out) <- fuseSuperScrema w inp_p form_p out_p inp_c form_c out_c   (new_inp, form') <-     fmap (second prunePreLambdaResults . toScrema) $
src/Futhark/Optimise/Fusion/TryFusion.hs view
@@ -30,7 +30,7 @@ import Futhark.IR.SOACS qualified as Futhark import Futhark.Optimise.Fusion.Composing import Futhark.Optimise.Fusion.Screma-import Futhark.Pass.ExtractKernels.ISRWIM (rwimPossible)+import Futhark.Transform.ISRWIM (rwimPossible) import Futhark.Transform.Rename (renameLambda) import Futhark.Transform.Substitute
src/Futhark/Optimise/InliningDeadFun.hs view
@@ -77,7 +77,14 @@             partition (noCallsTo to_inline_now . funDefName) funs        if null to_inline_now-        then pure (consts, funs)+        then+          -- Everything left calls something else we want to inline, so there is+          -- a cycle in the call graph. Break the cycle by dropping any function+          -- with more than one call site.+          let to_inline' = to_inline `S.intersection` calledOnce cg+           in if to_inline' == to_inline+                then pure (consts, funs)+                else recurse (i, vtable) (consts, funs) to_inline'         else do           let inlinemap =                 fdmap $ filter ((`S.member` to_inline_now) . funDefName) dont_inline_in@@ -259,6 +266,7 @@   bodyStms <$> inlineInBody fdmap (mkBody stms [])  inlineInBody ::+  forall m.   (MonadFreshNames m) =>   M.Map Name (FunDef SOACS) ->   Body SOACS ->@@ -286,15 +294,9 @@     inliner =       (identityMapper @SOACS)         { mapOnBody = const onBody,-          mapOnOp = onSOAC+          mapOnOp = traverseOpStms $ \_ -> inline . stmsToList         } -    onSOAC =-      mapSOACM identitySOACMapper {mapOnSOACLambda = onLambda}--    onLambda (Lambda params ret body) =-      Lambda params ret <$> onBody body- traceLocs :: Provenance -> StmAux () -> StmAux () traceLocs p aux =   aux {stmAuxLoc = stackProvenance p $ stmAuxLoc aux}@@ -321,7 +323,8 @@         runIdentity $           mapSOACM             identitySOACMapper-              { mapOnSOACLambda = pure . onLambda+              { mapOnSOACLambda = pure . onLambda,+                mapOnSOACExtLambda = pure . onLambda               }             soac       where@@ -354,7 +357,7 @@  -- | Remove functions not ultimately called from an entry point or a -- constant.-removeDeadFunctionsF :: Prog SOACS -> Prog SOACS+removeDeadFunctionsF :: (TraverseOpStms rep) => Prog rep -> Prog rep removeDeadFunctionsF prog =   let cg = buildCallGraph prog       live_funs = filter ((`isFunInCallGraph` cg) . funDefName) $ progFuns prog@@ -382,7 +385,7 @@  -- | @removeDeadFunctions prog@ removes the functions that are unreachable from -- the main function from the program.-removeDeadFunctions :: Pass SOACS SOACS+removeDeadFunctions :: (TraverseOpStms rep) => Pass rep rep removeDeadFunctions =   Pass     { passName = "Remove dead functions",
src/Futhark/Optimise/Simplify/Engine.hs view
@@ -1074,36 +1074,36 @@   simplify = traverse simplify  simplifyLambda ::-  (SimplifiableRep rep) =>+  (SimplifiableRep rep, Simplifiable t) =>   Names ->-  Lambda (Wise rep) ->-  SimpleM rep (Lambda (Wise rep), Stms (Wise rep))+  GLambda (Wise rep) t ->+  SimpleM rep (GLambda (Wise rep) t, Stms (Wise rep)) simplifyLambda extra_bound lam = do   par_blocker <- asksEngineEnv $ blockHoistPar . envHoistBlockers   simplifyLambdaMaybeHoist (par_blocker `orIf` hasFree extra_bound) mempty lam  simplifyLambdaNoHoisting ::-  (SimplifiableRep rep) =>-  Lambda (Wise rep) ->-  SimpleM rep (Lambda (Wise rep))+  (SimplifiableRep rep, Simplifiable t) =>+  GLambda (Wise rep) t ->+  SimpleM rep (GLambda (Wise rep) t) simplifyLambdaNoHoisting lam =   fst <$> simplifyLambdaMaybeHoist (isFalse False) mempty lam  simplifyLambdaMaybeHoist ::-  (SimplifiableRep rep) =>+  (SimplifiableRep rep, Simplifiable t) =>   BlockPred (Wise rep) ->   UT.UsageTable ->-  Lambda (Wise rep) ->-  SimpleM rep (Lambda (Wise rep), Stms (Wise rep))+  GLambda (Wise rep) t ->+  SimpleM rep (GLambda (Wise rep) t, Stms (Wise rep)) simplifyLambdaMaybeHoist = simplifyLambdaWith id  simplifyLambdaWith ::-  (SimplifiableRep rep) =>+  (SimplifiableRep rep, Simplifiable t) =>   (ST.SymbolTable (Wise rep) -> ST.SymbolTable (Wise rep)) ->   BlockPred (Wise rep) ->   UT.UsageTable ->-  Lambda (Wise rep) ->-  SimpleM rep (Lambda (Wise rep), Stms (Wise rep))+  GLambda (Wise rep) t ->+  SimpleM rep (GLambda (Wise rep) t, Stms (Wise rep)) simplifyLambdaWith f blocked usage lam@(Lambda params rettype body) = do   params' <- mapM (traverse simplify) params   let paramnames = namesFromList $ boundByLambda lam
src/Futhark/Optimise/Simplify/Rep.hs view
@@ -318,7 +318,7 @@ informBody (Body dec stms res) = mkWiseBody dec (informStms stms) res  -- | Construct a 'Wise' lambda.-informLambda :: (Informing rep) => Lambda rep -> Lambda (Wise rep)+informLambda :: (Informing rep) => GLambda rep t -> GLambda (Wise rep) t informLambda (Lambda ps ret body) = Lambda ps ret (informBody body)  -- | Construct a 'Wise' expression.
src/Futhark/Optimise/Simplify/Rules/BasicOp.hs view
@@ -186,7 +186,7 @@       False ruleBasicOp vtable pat aux (Update Unsafe dest is se)   | Just dest_t <- ST.lookupType dest vtable,-    isFullSlice (arrayShape dest_t) is = Simplify . auxing aux $+    isIdentitySlice (arrayShape dest_t) is = Simplify . auxing aux $       case se of         Var v | not $ null $ sliceDims is -> do           v_t <- lookupType v
src/Futhark/Optimise/Simplify/Rules/Index.hs view
@@ -101,6 +101,7 @@         not consuming,         not $ consumed arr,         Just (ordering, inds''') <- first concat . unzip <$> mapM okIdx inds'',+        length ordering == length matches,         Just perm <- L.sort ordering `isPermutationOf` ordering ->           if isIdentityPerm perm             then Just $ IndexResult cs arr . Slice <$> sequence inds'''@@ -333,3 +334,8 @@ -- -- In such cases we must actually insert a Rearrange operation to move -- the dimensions of the result appropriately.+--+-- Similarly, Every sliced dimension must correspond to exactly one dimension of+-- the result. If the indexing does not depend on one of them at all (e.g.+-- because the source array is a replicate), replacing the indexes with the+-- original slices would change the shape of the result.
src/Futhark/Optimise/Simplify/Rules/Match.hs view
@@ -5,13 +5,15 @@  import Control.Monad import Data.Either-import Data.List (partition, transpose, unzip4, zip5)+import Data.List (intersect, partition, tails, transpose, unzip4, zip5)+import Data.Map qualified as M import Futhark.Analysis.PrimExp.Convert import Futhark.Analysis.SymbolTable qualified as ST import Futhark.Analysis.UsageTable qualified as UT import Futhark.Construct import Futhark.IR import Futhark.Optimise.Simplify.Rule+import Futhark.Transform.Substitute import Futhark.Util  -- Does this case always match the scrutinees?@@ -142,7 +144,7 @@       (hoistings, (pes, ts, case_reses_tr, defbody_res')) =         (fmap unzip4 . partitionEithers) . map branchInvariant $           zip5 [0 ..] (patElems pat) ret (transpose case_reses) defbody_res-   in if null hoistings+   in if null hoistings || ifsort == MatchEquiv         then Skip         else Simplify $ do           ctx_fixes <- sequence hoistings@@ -207,6 +209,51 @@     reshapeResult se _ =       pure se +-- | Pairs @(i,j)@ with @i<j@ of the indexes of identical body results.+duplicateResults :: Body rep -> [(Int, Int)]+duplicateResults body = do+  (i, x) : rest <- tails $ zip [0 ..] $ bodyResult body+  (j, y) <- rest+  guard $ x == y+  pure (i, j)++-- | Combine duplicate branch results into a single result, with the name for+-- the duplicate bound after the branch. This is only valid when all branches+-- for the Match share the duplicate.+--+-- Example:+--+-- @+-- def f (b: bool) (xs: []i32) =+--   if b+--   then let ys = filter (> 0) xs+--        in (length ys, length ys)+--   else let zs = filter (< 0) xs+--        in (length zs, length zs)+-- @+unifyBranchDuplicate :: (BuilderOps rep) => TopDownRuleMatch rep+unifyBranchDuplicate _ pat aux (cond, cases, defbody, MatchDec ret ifsort)+  | defbody_dups <- duplicateResults defbody,+    cases_dups <- map (duplicateResults . caseBody) cases,+    -- We resolve only one duplicate per rule application. This is just to keep+    -- the logic simpler, although in principle we could resolve more at a time.+    (i, j) : _ <- foldl' intersect defbody_dups cases_dups = Simplify $ do+      let onBody (Body _ stms res) = mkBodyM stms $ without j res+          i_name = patNames pat !! i+          j_name = patNames pat !! j+          pat' =+            Pat . substituteNames (M.singleton j_name i_name) $+              without j (patElems pat)+          -- We need to adjust the existential references in the branch type.+          adjust = mapExt $ \x ->+            if x == j then i else if x >= j then x - 1 else x+      cases' <- mapM (traverse onBody) cases+      defbody' <- onBody defbody+      auxing aux . letBind pat' . Match cond cases' defbody' $+        MatchDec (map adjust $ without j ret) ifsort+      letBindNames [j_name] $ BasicOp $ SubExp $ Var i_name+  | otherwise = Skip+ -- | Remove the return values of a branch, that are not actually used -- after a branch.  Standard dead code removal can remove the branch -- if *none* of the return values are used, but this rule is more@@ -246,7 +293,8 @@ topDownRules :: (BuilderOps rep) => [TopDownRule rep] topDownRules =   [ RuleMatch ruleMatch,-    RuleMatch hoistBranchInvariant+    RuleMatch hoistBranchInvariant,+    RuleMatch unifyBranchDuplicate   ]  bottomUpRules :: (BuilderOps rep) => [BottomUpRule rep]
src/Futhark/Optimise/Sink.hs view
@@ -74,14 +74,15 @@ -- than 1. multiplicity :: (Constraints rep) => Stm rep -> M.Map VName Int multiplicity stm =-  case stmExp stm of-    Match cond cases defbody _ ->-      foldl' comb mempty $-        free 1 cond : free 1 defbody : map (free 1 . caseBody) cases-    Op {} -> free 2 stm-    Loop {} -> free 2 stm-    WithAcc {} -> free 2 stm-    _ -> free 1 stm+  free 1 (stmAux stm)+    <> case stmExp stm of+      Match cond cases defbody _ ->+        foldl' comb mempty $+          free 1 cond : free 1 defbody : map (free 1 . caseBody) cases+      Op {} -> free 2 stm+      Loop {} -> free 2 stm+      WithAcc {} -> free 2 stm+      _ -> free 1 stm   where     free k x = M.fromList $ map (,k) $ namesToList $ freeIn x     comb = M.unionWith (+)
src/Futhark/Pass/ExpandAllocations.hs view
@@ -25,12 +25,12 @@ import Futhark.Optimise.Simplify.Rep (addScopeWisdom) import Futhark.Pass import Futhark.Pass.ExplicitAllocations.GPU (explicitAllocationsInStms)-import Futhark.Pass.ExtractKernels.BlockedKernel (nonSegRed)-import Futhark.Pass.ExtractKernels.ToGPU (segThread)+import Futhark.Pass.Flatten.Builtins (mkSegSpace) import Futhark.Tools import Futhark.Transform.CopyPropagate (copyPropagateInFun)-import Futhark.Transform.Rename (renameStm)+import Futhark.Transform.Rename (renamePat, renameStm) import Futhark.Transform.Substitute+import Futhark.Transform.ToGPU (segThread) import Futhark.Util (mapAccumLM) import Prelude hiding (quot) @@ -157,7 +157,7 @@     onInput (shape, arrs, Just (op_lam, nes)) = do       bound_outside <- asks $ namesFromList . M.keys       let -- XXX: fake a SegLevel, which we don't have here.  We will not-          -- use it for anything, as we will not allow irregular+          -- use it for anything, as we will not allow nonuniform           -- allocations inside the update function.           lvl = SegThread SegNoVirt Nothing           (op_lam', lam_allocs) =@@ -171,7 +171,7 @@           throwError $             "Cannot handle un-sliceable allocation size: "               ++ prettyString v-              ++ "\nLikely cause: irregular nested operations inside accumulator update operator."+              ++ "\nLikely cause: nonuniform nested operations inside accumulator update operator."         [] ->           pure () @@ -234,7 +234,7 @@       throwError $         "Cannot handle un-sliceable allocation size: "           ++ prettyString v-          ++ "\nLikely cause: irregular nested operations inside parallel constructs."+          ++ "\nLikely cause: nonuniform nested operations inside parallel constructs."     Nothing ->       pure () @@ -953,6 +953,108 @@     addStms $ substituteNames substs stms   where     copy v = letExp (baseName v <> "_copy") $ BasicOp $ Replicate mempty $ Var v++data KernelInput = KernelInput+  { kernelInputName :: VName,+    kernelInputType :: Type,+    kernelInputArray :: VName,+    kernelInputIndices :: [SubExp]+  }+  deriving (Show)++readKernelInput :: (MonadBuilder m, Rep m ~ GPU.GPU) => KernelInput -> m ()+readKernelInput inp = do+  let pe = PatElem (kernelInputName inp) $ kernelInputType inp+  letBind (Pat [pe]) . BasicOp $+    case kernelInputType inp of+      Acc {} ->+        SubExp $ Var $ kernelInputArray inp+      _ ->+        Index (kernelInputArray inp) . Slice $+          map DimFix (kernelInputIndices inp)+            ++ map sliceDim (arrayDims (kernelInputType inp))++prepareRedOrScan ::+  (MonadBuilder m, Rep m ~ GPU.GPU) =>+  Certs ->+  SubExp ->+  Lambda GPU.GPU ->+  [VName] ->+  [(VName, SubExp)] ->+  [KernelInput] ->+  m (SegSpace, KernelBody GPU.GPU)+prepareRedOrScan cs w map_lam arrs ispace inps = do+  gtid <- newVName "gtid"+  space <- mkSegSpace $ ispace ++ [(gtid, w)]+  kbody <- fmap (uncurry (flip (Body ()))) $+    runBuilder $+      localScope (scopeOfSegSpace space) $ do+        mapM_ readKernelInput inps+        certifying cs . mapM_ readKernelInput $ do+          (p, arr) <- zip (lambdaParams map_lam) arrs+          pure $ KernelInput (paramName p) (paramType p) arr [Var gtid]+        res <- bodyBind (lambdaBody map_lam)+        forM res $ \(SubExpRes res_cs se) -> pure $ Returns ResultMaySimplify res_cs se++  pure (space, kbody)++segRed ::+  (MonadBuilder m, Rep m ~ GPU.GPU) =>+  SegOpLevel (Rep m) ->+  Pat (LetDec (Rep m)) ->+  Certs ->+  SubExp -> -- segment size+  [SegBinOp (Rep m)] ->+  Lambda (Rep m) ->+  [VName] ->+  [(VName, SubExp)] -> -- ispace = pair of (gtid, size) for the maps on "top" of this reduction+  [KernelInput] -> -- inps = inputs that can be looked up by using the gtids from ispace+  m (Stms (Rep m))+segRed lvl pat cs w ops map_lam arrs ispace inps = runBuilder_ $ do+  (kspace, kbody) <- prepareRedOrScan cs w map_lam arrs ispace inps+  letBind pat . Op . segOp $+    SegRed lvl kspace (lambdaReturnType map_lam) kbody ops++dummyDim ::+  (MonadBuilder m) =>+  Pat Type ->+  m (Pat Type, [(VName, SubExp)], m ())+dummyDim pat = do+  -- We add a unit-size segment on top to ensure that the result+  -- of the SegRed is an array, which we then immediately index.+  -- This is useful in the case that the value is used on the+  -- device afterwards, as this may save an expensive+  -- host-device copy (scalars are kept on the host, but arrays+  -- may be on the device).+  let addDummyDim t = t `arrayOfRow` intConst Int64 1+  pat' <- fmap addDummyDim <$> renamePat pat+  dummy <- newVName "dummy"+  let ispace = [(dummy, intConst Int64 1)]++  pure+    ( pat',+      ispace,+      forM_ (zip (patNames pat') (patNames pat)) $ \(from, to) -> do+        from_t <- lookupType from+        letBindNames [to] . BasicOp $+          case from_t of+            Acc {} -> SubExp $ Var from+            _ -> Index from $ fullSlice from_t [DimFix $ intConst Int64 0]+    )++nonSegRed ::+  (MonadBuilder m, Rep m ~ GPU.GPU) =>+  SegOpLevel (Rep m) ->+  Pat Type ->+  SubExp ->+  [SegBinOp (Rep m)] ->+  Lambda (Rep m) ->+  [VName] ->+  m (Stms (Rep m))+nonSegRed lvl pat w ops map_lam arrs = runBuilder_ $ do+  (pat', ispace, read_dummy) <- dummyDim pat+  addStms =<< segRed lvl pat' mempty w ops map_lam arrs ispace []+  read_dummy  -- Important for edge cases (#1838) that the Stms here still have the -- Allocs we are actually trying to get rid of.
src/Futhark/Pass/ExplicitAllocations.hs view
@@ -209,15 +209,16 @@ allocsForStm ::   (Allocable fromrep torep inner) =>   [Ident] ->+  StmAux a ->   Exp torep ->   AllocM fromrep torep (Stm torep)-allocsForStm idents e = do+allocsForStm idents aux e = do   def_space <- askDefaultSpace   hints <- expHints e   (rts, e') <- expReturns' e   pes <- allocsForPat def_space idents rts hints   dec <- mkExpDecM (Pat pes) e'-  pure $ Let (Pat pes) (defAux dec) e'+  pure $ Let (Pat pes) (aux {stmAuxDec = dec}) e'  patWithAllocations ::   (MonadBuilder m, Mem (Rep m) inner) =>@@ -726,7 +727,7 @@   where     allocInStms' [] = m     allocInStms' (stm : stms) = do-      allocstms <- collectStms_ $ auxing (stmAux stm) $ allocInStm stm+      allocstms <- collectStms_ $ allocInStm stm       addStms allocstms       let stms_consts = foldMap stmConsts allocstms           f env = env {envConsts = stms_consts <> envConsts env}@@ -736,8 +737,8 @@   (Allocable fromrep torep inner) =>   Stm fromrep ->   AllocM fromrep torep ()-allocInStm (Let (Pat pes) _ e) =-  addStm =<< allocsForStm (map patElemIdent pes) =<< allocInExp e+allocInStm (Let (Pat pes) aux e) =+  addStm =<< allocsForStm (map patElemIdent pes) aux =<< allocInExp e  allocInLambda ::   (Allocable fromrep torep inner) =>@@ -876,6 +877,7 @@ -- the idea, but which could perhaps be generalised. simplifyMatch ::   (Mem rep inner) =>+  MatchSort ->   [Case (Body rep)] ->   Body rep ->   [BranchTypeMem] ->@@ -883,7 +885,10 @@     Body rep,     [BranchTypeMem]   )-simplifyMatch cases defbody ts =+-- XXX: Do not simplify MatchEquivs+simplifyMatch MatchEquiv cases defbody ts =+  (cases, defbody, ts)+simplifyMatch _ cases defbody ts =   let case_reses = map (bodyResult . caseBody) cases       defbody_res = bodyResult defbody       (ctx_fixes, variant) =@@ -949,7 +954,7 @@   defbody'' <- addCtxToMatchBody reqs defbody'   cases'' <- mapM (traverse $ addCtxToMatchBody reqs) cases'   let (cases''', defbody''', rets') =-        simplifyMatch cases'' defbody'' $ mkBranchRet reqs+        simplifyMatch ifsort cases'' defbody'' $ mkBranchRet reqs   pure $ Match ses cases''' defbody''' $ MatchDec rets' ifsort   where     onCase (Case vs body) = first (Case vs) <$> allocInMatchBody rets body
− src/Futhark/Pass/ExtractKernels.hs 
@@ -1,892 +0,0 @@-{-# LANGUAGE LambdaCase #-}-{-# LANGUAGE TypeFamilies #-}---- | Kernel extraction.------ In the following, I will use the term "width" to denote the amount--- of immediate parallelism in a map - that is, the outer size of the--- array(s) being used as input.------ = Basic Idea------ If we have:------ @---   map---     map(f)---     stms_a...---     map(g)--- @------ Then we want to distribute to:------ @---   map---     map(f)---   map---     stms_a---   map---     map(g)--- @------ But for now only if------  (0) it can be done without creating irregular arrays.---      Specifically, the size of the arrays created by @map(f)@, by---      @map(g)@ and whatever is created by @stms_a@ that is also used---      in @map(g)@, must be invariant to the outermost loop.------  (1) the maps are _balanced_.  That is, the functions @f@ and @g@---      must do the same amount of work for every iteration.------ The advantage is that the map-nests containing @map(f)@ and--- @map(g)@ can now be trivially flattened at no cost, thus exposing--- more parallelism.  Note that the @stms_a@ map constitutes array--- expansion, which requires additional storage.------ = Distributing Sequential Loops------ As a starting point, sequential loops are treated like scalar--- expressions.  That is, not distributed.  However, sometimes it can--- be worthwhile to distribute if they contain a map:------ @---   map---     loop---       map---     map--- @------ If we distribute the loop and interchange the outer map into the--- loop, we get this:------ @---   loop---     map---       map---   map---     map--- @------ Now more parallelism may be available.------ = Unbalanced Maps------ Unbalanced maps will as a rule be sequentialised, but sometimes,--- there is another way.  Assume we find this:------ @---   map---     map(f)---       map(g)---     map--- @------ Presume that @map(f)@ is unbalanced.  By the simple rule above, we--- would then fully sequentialise it, resulting in this:------ @---   map---     loop---   map---     map--- @------ == Balancing by Loop Interchange------ The above is not ideal, as we cannot flatten the @map-loop@ nest,--- and we are thus limited in the amount of parallelism available.------ But assume now that the width of @map(g)@ is invariant to the outer--- loop.  Then if possible, we can interchange @map(f)@ and @map(g)@,--- sequentialise @map(f)@ and distribute, interchanging the outer--- parallel loop into the sequential loop:------ @---   loop(f)---     map---       map(g)---   map---     map--- @------ After flattening the two nests we can obtain more parallelism.------ When distributing a map, we also need to distribute everything that--- the map depends on - possibly as its own map.  When distributing a--- set of scalar bindings, we will need to know which of the binding--- results are used afterwards.  Hence, we will need to compute usage--- information.------ = Redomap------ Redomap can be handled much like map.  Distributed loops are--- distributed as maps, with the parameters corresponding to the--- neutral elements added to their bodies.  The remaining loop will--- remain a redomap.  Example:------ @--- redomap(op,---         fn (v) =>---           map(f)---           map(g),---         e,a)--- @------ distributes to------ @--- let b = map(fn v =>---               let acc = e---               map(f),---               a)--- redomap(op,---         fn (v,dist) =>---           map(g),---         e,a,b)--- @------ Note that there may be further kernel extraction opportunities--- inside the @map(f)@.  The downside of this approach is that the--- intermediate array (@b@ above) must be written to main memory.  An--- often better approach is to just turn the entire @redomap@ into a--- single kernel.-module Futhark.Pass.ExtractKernels (extractKernels) where--import Control.Monad-import Control.Monad.RWS.Strict-import Control.Monad.Reader-import Data.Bifunctor (first)-import Data.Maybe-import Futhark.IR.GPU-import Futhark.IR.SOACS-import Futhark.IR.SOACS.Simplify (simplifyStms)-import Futhark.MonadFreshNames-import Futhark.Pass-import Futhark.Pass.ExtractKernels.BlockedKernel-import Futhark.Pass.ExtractKernels.DistributeNests-import Futhark.Pass.ExtractKernels.Distribution-import Futhark.Pass.ExtractKernels.ISRWIM-import Futhark.Pass.ExtractKernels.Intrablock-import Futhark.Pass.ExtractKernels.StreamKernel-import Futhark.Pass.ExtractKernels.ToGPU-import Futhark.Tools-import Futhark.Transform.FirstOrderTransform qualified as FOT-import Futhark.Transform.Rename-import Futhark.Util.Log-import Prelude hiding (log)---- | Transform a program using SOACs to a program using explicit--- kernels, using the kernel extraction transformation.-extractKernels :: Pass SOACS GPU-extractKernels =-  Pass-    { passName = "extract kernels",-      passDescription = "Perform kernel extraction",-      passFunction = transformProg-    }--transformProg :: Prog SOACS -> PassM (Prog GPU)-transformProg prog = do-  consts' <- runDistribM $ transformStms mempty $ stmsToList $ progConsts prog-  funs' <- mapM (transformFunDef $ scopeOf consts') $ progFuns prog-  pure $-    prog-      { progConsts = consts',-        progFuns = funs'-      }---- In order to generate more stable threshold names, we keep track of--- the numbers used for thresholds separately from the ordinary name--- source,-data State = State-  { stateNameSource :: VNameSource,-    stateThresholdCounter :: Int-  }--newtype DistribM a = DistribM (RWS (Scope GPU) Log State a)-  deriving-    ( Functor,-      Applicative,-      Monad,-      HasScope GPU,-      LocalScope GPU,-      MonadState State,-      MonadLogger-    )--instance MonadFreshNames DistribM where-  getNameSource = gets stateNameSource-  putNameSource src = modify $ \s -> s {stateNameSource = src}--runDistribM ::-  (MonadLogger m, MonadFreshNames m) =>-  DistribM a ->-  m a-runDistribM (DistribM m) = do-  (x, msgs) <- modifyNameSource $ \src ->-    let (x, s, msgs) = runRWS m mempty (State src 0)-     in ((x, msgs), stateNameSource s)-  addLog msgs-  pure x--transformFunDef ::-  (MonadFreshNames m, MonadLogger m) =>-  Scope GPU ->-  FunDef SOACS ->-  m (FunDef GPU)-transformFunDef scope (FunDef entry attrs name rettype params body) = runDistribM $ do-  body' <--    localScope (scope <> scopeOfFParams params) $-      transformBody mempty body-  pure $ FunDef entry attrs name rettype params body'--type GPUStms = Stms GPU--transformBody :: KernelPath -> Body SOACS -> DistribM (Body GPU)-transformBody path body = do-  stms <- transformStms path $ stmsToList $ bodyStms body-  pure $ mkBody stms $ bodyResult body--transformStms :: KernelPath -> [Stm SOACS] -> DistribM GPUStms-transformStms _ [] =-  pure mempty-transformStms path (stm : stms) =-  sequentialisedUnbalancedStm stm >>= \case-    Nothing -> do-      stm' <- transformStm path stm-      inScopeOf stm' $-        (stm' <>) <$> transformStms path stms-    Just stms' ->-      transformStms path $ stmsToList stms' <> stms--unbalancedLambda :: Lambda SOACS -> Bool-unbalancedLambda orig_lam =-  unbalancedBody (namesFromList $ map paramName $ lambdaParams orig_lam) $-    lambdaBody orig_lam-  where-    subExpBound (Var i) bound = i `nameIn` bound-    subExpBound (Constant _) _ = False--    unbalancedBody bound body =-      any (unbalancedStm (bound <> boundInBody body) . stmExp) $-        bodyStms body--    -- XXX - our notion of balancing is probably still too naive.-    unbalancedStm bound (Op (Stream w _ _ _)) =-      w `subExpBound` bound-    unbalancedStm bound (Op (Screma w _ _)) =-      w `subExpBound` bound-    unbalancedStm _ Op {} =-      False-    unbalancedStm _ Loop {} = False-    unbalancedStm bound (WithAcc _ lam) =-      unbalancedBody bound (lambdaBody lam)-    unbalancedStm bound (Match ses cases defbody _) =-      any (`subExpBound` bound) ses-        && ( any (unbalancedBody bound . caseBody) cases-               || unbalancedBody bound defbody-           )-    unbalancedStm _ (BasicOp _) =-      False-    unbalancedStm _ Apply {} = False--sequentialisedUnbalancedStm :: Stm SOACS -> DistribM (Maybe (Stms SOACS))-sequentialisedUnbalancedStm (Let pat _ (Op soac@(Screma _ _ form)))-  | Just (_, lam2) <- isRedomapSOAC form,-    unbalancedLambda lam2,-    lambdaContainsParallelism lam2 = do-      types <- asksScope scopeForSOACs-      Just . snd <$> runBuilderT (FOT.transformSOAC pat soac) types-sequentialisedUnbalancedStm _ =-  pure Nothing--cmpSizeLe ::-  Name ->-  SizeClass ->-  [SubExp] ->-  DistribM ((SubExp, Name), Stms GPU)-cmpSizeLe desc size_class to_what = do-  x <- gets stateThresholdCounter-  modify $ \s -> s {stateThresholdCounter = x + 1}-  let size_key = desc <> "_" <> nameFromString (show x)-  runBuilder $ do-    to_what' <--      letSubExp "comparatee"-        =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) to_what-    cmp_res <- letSubExp desc $ Op $ SizeOp $ CmpSizeLe size_key size_class to_what'-    pure (cmp_res, size_key)--kernelAlternatives ::-  (MonadFreshNames m, HasScope GPU m) =>-  Pat Type ->-  Body GPU ->-  [(SubExp, Body GPU)] ->-  m (Stms GPU)-kernelAlternatives pat default_body [] = runBuilder_ $ do-  ses <- bodyBind default_body-  forM_ (zip (patNames pat) ses) $ \(name, SubExpRes cs se) ->-    certifying cs $ letBindNames [name] $ BasicOp $ SubExp se-kernelAlternatives pat default_body ((cond, alt) : alts) = runBuilder_ $ do-  alts_pat <- fmap Pat . forM (patElems pat) $ \pe -> do-    name <- newName $ patElemName pe-    pure pe {patElemName = name}--  alt_stms <- kernelAlternatives alts_pat default_body alts-  let alt_body = mkBody alt_stms $ varsRes $ patNames alts_pat--  letBind pat . Match [cond] [Case [Just $ BoolValue True] alt] alt_body $-    MatchDec (staticShapes (patTypes pat)) MatchEquiv--transformLambda :: KernelPath -> Lambda SOACS -> DistribM (Lambda GPU)-transformLambda path (Lambda params ret body) =-  Lambda params ret-    <$> localScope (scopeOfLParams params) (transformBody path body)--versionScanRed ::-  KernelPath ->-  Pat Type ->-  StmAux () ->-  SubExp ->-  Lambda SOACS ->-  DistribM (Stms GPU) ->-  DistribM (Body GPU) ->-  ([(Name, Bool)] -> DistribM (Body GPU)) ->-  DistribM (Stms GPU)-versionScanRed path pat aux w map_lam paralleliseOuter outerParallelBody innerParallelBody =-  if not (lambdaContainsParallelism map_lam)-    || ("sequential_inner" `inAttrs` stmAuxAttrs aux)-    then paralleliseOuter-    else do-      ((outer_suff, outer_suff_key), suff_stms) <--        sufficientParallelism "suff_outer_screma" [w] path Nothing--      outer_stms <- outerParallelBody-      inner_stms <- innerParallelBody ((outer_suff_key, False) : path)--      (suff_stms <>) <$> kernelAlternatives pat inner_stms [(outer_suff, outer_stms)]--transformStm :: KernelPath -> Stm SOACS -> DistribM GPUStms-transformStm _ stm-  | "sequential" `inAttrs` stmAuxAttrs (stmAux stm) =-      runBuilder_ $ FOT.transformStmRecursively stm-transformStm path (Let pat aux (Op soac))-  | "sequential_outer" `inAttrs` stmAuxAttrs aux =-      transformStms path . stmsToList . fmap (certify (stmAuxCerts aux))-        =<< runBuilder_ (FOT.transformSOAC pat soac)-transformStm path (Let pat aux (Match c cases defbody rt)) = do-  cases' <- mapM (traverse $ transformBody path) cases-  defbody' <- transformBody path defbody-  pure $ oneStm $ Let pat aux $ Match c cases' defbody' rt-transformStm path (Let pat aux (WithAcc inputs lam)) =-  oneStm . Let pat aux-    <$> (WithAcc (map transformInput inputs) <$> transformLambda path lam)-  where-    transformInput (shape, arrs, op) =-      (shape, arrs, fmap (first soacsLambdaToGPU) op)-transformStm path (Let pat aux (Loop merge form body)) =-  localScope (scopeOfLoopForm form <> scopeOfFParams params) $-    oneStm . Let pat aux . Loop merge form <$> transformBody path body-  where-    params = map fst merge-transformStm path (Let pat aux (Op (Screma w arrs form)))-  | Just lam <- isMapSOAC form =-      onMap path $ MapLoop pat aux w lam arrs-transformStm path (Let pat aux (Op (Screma w arrs form)))-  | Just scans <- isScanSOAC form,-    Scan scan_lam nes <- singleScan scans,-    Just do_iswim <- iswim pat w scan_lam $ zip nes arrs = do-      types <- asksScope scopeForSOACs-      transformStms path . stmsToList . snd =<< runBuilderT (certifying cs do_iswim) types-  | Just (post_lam, scans, map_lam) <- isMaposcanomapSOAC form = do-      let paralleliseOuter = runBuilder_ $ do-            scan_ops <- forM scans $ \(Scan scan_lam nes) -> do-              (scan_lam', nes', shape) <- determineReduceOp scan_lam nes-              let scan_lam'' = soacsLambdaToGPU scan_lam'-              pure $ SegBinOp Noncommutative scan_lam'' nes' shape-            let map_lam_sequential = soacsLambdaToGPU map_lam-                post_op = SegPostOp $ soacsLambdaToGPU post_lam-            lvl <- segThreadCapped [w] "segscan" $ NoRecommendation SegNoVirt-            addStms . fmap (certify cs)-              =<< segScan lvl pat mempty w scan_ops map_lam_sequential post_op arrs [] []--          outerParallelBody =-            renameBody-              =<< (mkBody <$> paralleliseOuter <*> pure (varsRes (patNames pat)))--          paralleliseInner path' = do-            (mapstm, scanstm, poststm) <--              maposcanomapToMapScanAndMap pat (w, post_lam, scans, map_lam, arrs)-            types <- asksScope scopeForSOACs-            transformStms path' . stmsToList <=< (`runBuilderT_` types) $-              addStms-                =<< simplifyStms-                  ( stmsFromList-                      [ certify cs mapstm,-                        certify cs scanstm,-                        certify cs poststm-                      ]-                  )--          innerParallelBody path' =-            renameBody-              =<< (mkBody <$> paralleliseInner path' <*> pure (varsRes (patNames pat)))--      versionScanRed path pat aux w map_lam paralleliseOuter outerParallelBody innerParallelBody-  where-    cs = stmAuxCerts aux-transformStm path (Let res_pat aux (Op (Screma w arrs form)))-  | Just [Reduce comm red_fun nes] <- isReduceSOAC form,-    let comm'-          | commutativeLambda red_fun = Commutative-          | otherwise = comm,-    Just do_irwim <- irwim res_pat w comm' red_fun $ zip nes arrs = do-      types <- asksScope scopeForSOACs-      stms <- fst <$> runBuilderT (simplifyStms =<< collectStms_ (auxing aux do_irwim)) types-      transformStms path $ stmsToList stms-transformStm path (Let pat aux (Op (Screma w arrs form)))-  | Just (reds, map_lam) <- isRedomapSOAC form = do-      let paralleliseOuter = runBuilder_ $ do-            red_ops <- forM reds $ \(Reduce comm red_lam nes) -> do-              (red_lam', nes', shape) <- determineReduceOp red_lam nes-              let comm'-                    | commutativeLambda red_lam' = Commutative-                    | otherwise = comm-                  red_lam'' = soacsLambdaToGPU red_lam'-              pure $ SegBinOp comm' red_lam'' nes' shape-            let map_lam_sequential = soacsLambdaToGPU map_lam-            lvl <- segThreadCapped [w] "segred" $ NoRecommendation SegNoVirt-            addStms . fmap (certify cs)-              =<< nonSegRed lvl pat w red_ops map_lam_sequential arrs--          outerParallelBody =-            renameBody-              =<< (mkBody <$> paralleliseOuter <*> pure (varsRes (patNames pat)))--          paralleliseInner path' = do-            (mapstm, redstm) <--              redomapToMapAndReduce pat (w, reds, map_lam, arrs)-            types <- asksScope scopeForSOACs-            transformStms path' . stmsToList <=< (`runBuilderT_` types) $-              addStms =<< simplifyStms (stmsFromList [certify cs mapstm, certify cs redstm])--          innerParallelBody path' =-            renameBody-              =<< (mkBody <$> paralleliseInner path' <*> pure (varsRes (patNames pat)))--      versionScanRed path pat aux w map_lam paralleliseOuter outerParallelBody innerParallelBody-  where-    cs = stmAuxCerts aux-transformStm path (Let pat aux (Op (Screma w arrs form))) = do-  -- This screma is too complicated for us to immediately do-  -- anything, so split it up and try again.-  scope <- asksScope scopeForSOACs-  transformStms path . map (certify (stmAuxCerts aux)) . stmsToList . snd-    =<< runBuilderT (dissectScrema pat w form arrs) scope-transformStm path (Let pat _ (Op (Stream w arrs nes fold_fun))) = do-  -- Remove the stream and leave the body parallel.  It will be-  -- distributed.-  types <- asksScope scopeForSOACs-  transformStms path . stmsToList . snd-    =<< runBuilderT (sequentialStreamWholeArray pat w nes fold_fun arrs) types-transformStm _ (Let orig_pat aux (Op (Hist w imgs ops bucket_fun))) = do-  let bfun' = soacsLambdaToGPU bucket_fun--  -- It is important not to launch unnecessarily many threads for-  -- histograms, because it may mean we unnecessarily need to reduce-  -- subhistograms as well.-  runBuilder_ $ do-    lvl <- segThreadCapped [w] "seghist" $ NoRecommendation SegNoVirt-    addStms =<< histKernel onLambda lvl orig_pat [] [] (stmAuxCerts aux) w ops bfun' imgs-  where-    onLambda = pure . soacsLambdaToGPU-transformStm _ stm =-  runBuilder_ $ FOT.transformStmRecursively stm--sufficientParallelism ::-  Name ->-  [SubExp] ->-  KernelPath ->-  Maybe Int64 ->-  DistribM ((SubExp, Name), Stms GPU)-sufficientParallelism desc ws path def =-  cmpSizeLe desc (SizeThreshold path def) ws---- | Intra-group parallelism is worthwhile if the lambda contains more--- than one instance of non-map nested parallelism, or any nested--- parallelism inside a loop.-worthIntrablock :: Lambda SOACS -> Bool-worthIntrablock lam = bodyInterest (lambdaBody lam) > 1-  where-    bodyInterest body =-      sum $ interest <$> bodyStms body-    interest stm-      | "sequential" `inAttrs` attrs =-          0 :: Int-      | Op (Screma w _ form) <- stmExp stm,-        Just lam' <- isMapSOAC form =-          mapLike w lam'-      | Loop _ _ body <- stmExp stm =-          bodyInterest body * 10-      | Match _ cases defbody _ <- stmExp stm =-          foldl-            max-            (bodyInterest defbody)-            (map (bodyInterest . caseBody) cases)-      | Op (Screma w _ (ScremaForm lam' _ _ _)) <- stmExp stm =-          zeroIfTooSmall w + bodyInterest (lambdaBody lam')-      | Op (Stream _ _ _ lam') <- stmExp stm =-          bodyInterest $ lambdaBody lam'-      | WithAcc _ lam' <- stmExp stm =-          bodyInterest $ lambdaBody lam'-      | otherwise =-          0-      where-        attrs = stmAuxAttrs $ stmAux stm-        sequential_inner = "sequential_inner" `inAttrs` attrs--        zeroIfTooSmall (Constant (IntValue x))-          | intToInt64 x < 32 = 0-        zeroIfTooSmall _ = 1--        mapLike w lam' =-          if sequential_inner-            then 0-            else max (zeroIfTooSmall w) (bodyInterest (lambdaBody lam'))---- | A lambda is worth sequentialising if it contains enough nested--- parallelism of an interesting kind.-worthSequentialising :: Lambda SOACS -> Bool-worthSequentialising lam = bodyInterest (0 :: Int) (lambdaBody lam) > 1-  where-    bodyInterest depth body =-      sum $ interest depth <$> bodyStms body-    interest depth stm-      | "sequential" `inAttrs` attrs =-          0 :: Int-      | Op (Screma _ _ form@(ScremaForm lam' _ _ _)) <- stmExp stm,-        isJust $ isMapSOAC form =-          if sequential_inner-            then 0-            else bodyInterest (depth + 1) (lambdaBody lam')-      | Loop _ ForLoop {} body <- stmExp stm =-          bodyInterest (depth + 1) body * 10-      | WithAcc _ withacc_lam <- stmExp stm =-          bodyInterest (depth + 1) (lambdaBody withacc_lam)-      | Op (Screma _ _ form@(ScremaForm lam' _ _ _)) <- stmExp stm =-          1-            + bodyInterest (depth + 1) (lambdaBody lam')-            +-            -- Give this a bigger score if it's a redomap just inside-            -- the the outer lambda, as these are often tileable and-            -- thus benefit more from sequentialisation.-            case (isRedomapSOAC form, depth) of-              (Just _, 0) -> 1-              _ -> 0-      | otherwise =-          0-      where-        attrs = stmAuxAttrs $ stmAux stm-        sequential_inner = "sequential_inner" `inAttrs` attrs--onTopLevelStms ::-  KernelPath ->-  Stms SOACS ->-  DistNestT GPU DistribM GPUStms-onTopLevelStms path stms =-  liftInner $ transformStms path $ stmsToList stms--onMap :: KernelPath -> MapLoop -> DistribM GPUStms-onMap path (MapLoop pat aux w lam arrs) = do-  types <- askScope-  let loopnest = MapNesting pat aux w $ zip (lambdaParams lam) arrs-      env path' =-        DistEnv-          { distNest = singleNesting (Nesting mempty loopnest),-            distScope =-              scopeOfPat pat-                <> scopeForGPU (scopeOf lam)-                <> types,-            distOnInnerMap = onInnerMap path',-            distOnTopLevelStms = onTopLevelStms path',-            distSegLevel = segThreadCapped,-            distOnSOACSStms = pure . oneStm . soacsStmToGPU,-            distOnSOACSLambda = pure . soacsLambdaToGPU-          }-      exploitInnerParallelism path' =-        runDistNestT (env path') $-          distributeMapBodyStms acc (bodyStms $ lambdaBody lam)--  let exploitOuterParallelism path' = do-        let lam' = soacsLambdaToGPU lam-        runDistNestT (env path') $-          distribute $-            addStmsToAcc (bodyStms $ lambdaBody lam') acc--  onMap' (newKernel loopnest) path exploitOuterParallelism exploitInnerParallelism pat lam-  where-    acc =-      DistAcc-        { distTargets = singleTarget (pat, bodyResult $ lambdaBody lam),-          distStms = mempty-        }--onlyExploitIntra :: Attrs -> Bool-onlyExploitIntra attrs =-  AttrComp "incremental_flattening" ["only_intra"] `inAttrs` attrs--mayExploitOuter :: Attrs -> Bool-mayExploitOuter attrs =-  not $-    AttrComp "incremental_flattening" ["no_outer"]-      `inAttrs` attrs-      || AttrComp "incremental_flattening" ["only_inner"]-        `inAttrs` attrs--mayExploitIntra :: Attrs -> Bool-mayExploitIntra attrs =-  not $-    AttrComp "incremental_flattening" ["no_intra"]-      `inAttrs` attrs-      || AttrComp "incremental_flattening" ["only_inner"]-        `inAttrs` attrs---- The minimum amount of inner parallelism we require (by default) in--- intra-group versions.  Less than this is usually pointless on a GPU--- (but we allow tuning to change it).-intraMinInnerPar :: Int64-intraMinInnerPar = 32 -- One NVIDIA warp--onMap' ::-  KernelNest ->-  KernelPath ->-  (KernelPath -> DistribM (Stms GPU)) ->-  (KernelPath -> DistribM (Stms GPU)) ->-  Pat Type ->-  Lambda SOACS ->-  DistribM (Stms GPU)-onMap' loopnest path mk_seq_stms mk_par_stms pat lam = do-  -- Some of the control flow here looks a bit convoluted because we-  -- are trying to avoid generating unneeded threshold parameters,-  -- which means we need to do all the pruning checks up front.--  types <- askScope--  let only_intra = onlyExploitIntra (stmAuxAttrs aux)-      may_intra = worthIntrablock lam && mayExploitIntra attrs--  intra <--    if only_intra || may_intra-      then flip runReaderT types $ intrablockParallelise loopnest lam-      else pure Nothing--  case intra of-    _ | "sequential_inner" `inAttrs` attrs -> do-      seq_body <- renameBody =<< mkBody <$> mk_seq_stms path <*> pure res-      kernelAlternatives pat seq_body []-    ---    Nothing-      | not only_intra,-        Just m <- mkSeqAlts -> do-          (outer_suff, outer_suff_key, outer_suff_stms, seq_body) <- m-          par_body <--            renameBody-              =<< mkBody-                <$> mk_par_stms ((outer_suff_key, False) : path)-                <*> pure res-          (outer_suff_stms <>) <$> kernelAlternatives pat par_body [(outer_suff, seq_body)]-      ---      | otherwise -> do-          par_body <- renameBody =<< mkBody <$> mk_par_stms path <*> pure res-          kernelAlternatives pat par_body []-    ---    Just intra'@(_, _, log, intra_prelude, intra_stms)-      | only_intra -> do-          addLog log-          group_par_body <- renameBody $ mkBody intra_stms res-          (intra_prelude <>) <$> kernelAlternatives pat group_par_body []-      ---      | otherwise -> do-          addLog log--          case mkSeqAlts of-            Nothing -> do-              (group_par_body, intra_ok, intra_suff_key, intra_suff_stms) <--                checkSuffIntraPar path intra'--              par_body <--                renameBody-                  =<< mkBody-                    <$> mk_par_stms ((intra_suff_key, False) : path)-                    <*> pure res--              (intra_suff_stms <>)-                <$> kernelAlternatives pat par_body [(intra_ok, group_par_body)]-            Just m -> do-              (outer_suff, outer_suff_key, outer_suff_stms, seq_body) <- m--              (group_par_body, intra_ok, intra_suff_key, intra_suff_stms) <--                checkSuffIntraPar ((outer_suff_key, False) : path) intra'--              par_body <--                renameBody-                  =<< mkBody-                    <$> mk_par_stms-                      ( [ (outer_suff_key, False),-                          (intra_suff_key, False)-                        ]-                          ++ path-                      )-                    <*> pure res--              ((outer_suff_stms <> intra_suff_stms) <>)-                <$> kernelAlternatives-                  pat-                  par_body-                  [(outer_suff, seq_body), (intra_ok, group_par_body)]-  where-    nest_ws = kernelNestWidths loopnest-    res = varsRes $ patNames pat-    aux = loopNestingAux $ innermostKernelNesting loopnest-    attrs = stmAuxAttrs aux--    mkSeqAlts-      | worthSequentialising lam,-        mayExploitOuter attrs = Just $ do-          ((outer_suff, outer_suff_key), outer_suff_stms) <- checkSuffOuterPar-          seq_body <--            renameBody-              =<< mkBody-                <$> mk_seq_stms ((outer_suff_key, True) : path)-                <*> pure res-          pure (outer_suff, outer_suff_key, outer_suff_stms, seq_body)-      | otherwise =-          Nothing--    checkSuffOuterPar =-      sufficientParallelism "suff_outer_par" nest_ws path Nothing--    checkSuffIntraPar-      path'-      ((_intra_min_par, intra_avail_par), tblock_size, _, intra_prelude, intra_stms) = do-        -- We must check that all intra-group parallelism fits in a group.-        ((intra_ok, intra_suff_key), intra_suff_stms) <- do-          ((intra_suff, suff_key), check_suff_stms) <--            sufficientParallelism-              "suff_intra_par"-              [intra_avail_par]-              path'-              (Just intraMinInnerPar)--          runBuilder $ do-            addStms intra_prelude--            max_tblock_size <--              letSubExp "max_tblock_size" $ Op $ SizeOp $ GetSizeMax SizeThreadBlock-            fits <--              letSubExp "fits" $-                BasicOp $-                  CmpOp (CmpSle Int64) tblock_size max_tblock_size--            addStms check_suff_stms--            intra_ok <- letSubExp "intra_suff_and_fits" $ BasicOp $ BinOp LogAnd fits intra_suff-            pure (intra_ok, suff_key)--        group_par_body <- renameBody $ mkBody intra_stms res-        pure (group_par_body, intra_ok, intra_suff_key, intra_suff_stms)--removeUnusedMapResults ::-  Pat Type ->-  [SubExpRes] ->-  Lambda rep ->-  Maybe ([Int], Pat Type, Lambda rep)-removeUnusedMapResults (Pat pes) res lam = do-  let (pes', body_res) =-        unzip $ filter (used . fst) $ zip pes $ bodyResult (lambdaBody lam)-  perm <- map (Var . patElemName) pes' `isPermutationOf` map resSubExp res-  pure (perm, Pat pes', lam {lambdaBody = (lambdaBody lam) {bodyResult = body_res}})-  where-    used pe = patElemName pe `nameIn` freeIn res--onInnerMap ::-  KernelPath ->-  MapLoop ->-  DistAcc GPU ->-  DistNestT GPU DistribM (DistAcc GPU)-onInnerMap path maploop@(MapLoop pat aux w lam arrs) acc-  | unbalancedLambda lam,-    lambdaContainsParallelism lam =-      flip addStmToAcc acc =<< mapLoopStm maploop-  | otherwise =-      (distributeSingleStm acc =<< mapLoopStm maploop) >>= \case-        Just (post_kernels, res, nest, acc')-          | Just (perm, pat', lam') <- removeUnusedMapResults pat res lam -> do-              addPostStms post_kernels-              multiVersion perm nest acc' pat' lam'-        _ -> distributeMap maploop acc-  where-    discardTargets acc' =-      -- FIXME: work around bogus targets.-      acc' {distTargets = singleTarget (mempty, mempty)}--    -- GHC 9.2 loops without the type annotation.-    generate ::-      [Int] ->-      KernelNest ->-      Pat Type ->-      Lambda SOACS ->-      DistEnv GPU DistribM ->-      Scope GPU ->-      DistribM (Stms GPU)-    generate perm nest pat' lam' dist_env extra_scope = localScope extra_scope $ do-      let maploop' = MapLoop pat' aux w lam' arrs--          exploitInnerParallelism path' = do-            let dist_env' =-                  dist_env-                    { distOnTopLevelStms = onTopLevelStms path',-                      distOnInnerMap = onInnerMap path'-                    }-            runDistNestT dist_env' . inNesting nest . localScope extra_scope $-              discardTargets-                <$> distributeMap maploop' acc {distStms = mempty}-      -- Normally the permutation is for the output pattern, but-      -- we can't really change that, so we change the result-      -- order instead.-      let lam_res' =-            rearrangeShape (rearrangeInverse perm) $-              bodyResult $-                lambdaBody lam'-          lam'' = lam' {lambdaBody = (lambdaBody lam') {bodyResult = lam_res'}}-          map_nesting = MapNesting pat' aux w $ zip (lambdaParams lam') arrs-          nest' = pushInnerKernelNesting (pat', lam_res') map_nesting nest--      -- XXX: we do not construct a new KernelPath when-      -- sequentialising.  This is only OK as long as further-      -- versioning does not take place down that branch (it currently-      -- does not).-      (sequentialised_kernel, nestw_stms) <- localScope extra_scope $ do-        let sequentialised_lam = soacsLambdaToGPU lam''-        constructKernel segThreadCapped nest' $ lambdaBody sequentialised_lam--      let outer_pat = loopNestingPat $ fst nest-      (nestw_stms <>)-        <$> onMap'-          nest'-          path-          (const $ pure $ oneStm sequentialised_kernel)-          exploitInnerParallelism-          outer_pat-          lam''--    multiVersion perm nest acc' pat' lam' = do-      -- The kernel can be distributed by itself, so now we can-      -- decide whether to just sequentialise, or exploit inner-      -- parallelism.-      dist_env <- ask-      let extra_scope = targetsScope $ distTargets acc'--      stms <- liftInner $ generate perm nest pat' lam' dist_env extra_scope-      postStm stms-      pure acc'
− src/Futhark/Pass/ExtractKernels/BlockedKernel.hs 
@@ -1,257 +0,0 @@-{-# LANGUAGE TypeFamilies #-}--module Futhark.Pass.ExtractKernels.BlockedKernel-  ( DistRep,-    MkSegLevel,-    ThreadRecommendation (..),-    segRed,-    nonSegRed,-    segScan,-    segHist,-    segMap,-    mapKernel,-    KernelInput (..),-    readKernelInput,-    mkSegSpace,-    dummyDim,-  )-where--import Control.Monad-import Futhark.Analysis.PrimExp-import Futhark.IR-import Futhark.IR.Aliases (AliasableRep)-import Futhark.IR.GPU.Op (SegVirt (..))-import Futhark.IR.SegOp-import Futhark.MonadFreshNames-import Futhark.Tools-import Futhark.Transform.Rename-import Prelude hiding (quot)---- | Constraints pertinent to performing distribution/flattening.-type DistRep rep =-  ( Buildable rep,-    HasSegOp rep,-    BuilderOps rep,-    LetDec rep ~ Type,-    ExpDec rep ~ (),-    BodyDec rep ~ (),-    AliasableRep rep-  )--data ThreadRecommendation = ManyThreads | NoRecommendation SegVirt--type MkSegLevel rep m =-  [SubExp] -> Name -> ThreadRecommendation -> BuilderT rep m (SegOpLevel rep)--mkSegSpace :: (MonadFreshNames m) => [(VName, SubExp)] -> m SegSpace-mkSegSpace dims = SegSpace <$> newVName "phys_tid" <*> pure dims--prepareRedOrScan ::-  (MonadBuilder m, DistRep (Rep m)) =>-  Certs ->-  SubExp ->-  Lambda (Rep m) ->-  [VName] ->-  [(VName, SubExp)] ->-  [KernelInput] ->-  m (SegSpace, KernelBody (Rep m))-prepareRedOrScan cs w map_lam arrs ispace inps = do-  gtid <- newVName "gtid"-  space <- mkSegSpace $ ispace ++ [(gtid, w)]-  kbody <- fmap (uncurry (flip (Body ()))) $-    runBuilder $-      localScope (scopeOfSegSpace space) $ do-        mapM_ readKernelInput inps-        certifying cs . mapM_ readKernelInput $ do-          (p, arr) <- zip (lambdaParams map_lam) arrs-          pure $ KernelInput (paramName p) (paramType p) arr [Var gtid]-        res <- bodyBind (lambdaBody map_lam)-        forM res $ \(SubExpRes res_cs se) -> pure $ Returns ResultMaySimplify res_cs se--  pure (space, kbody)--segRed ::-  (MonadFreshNames m, DistRep rep, HasScope rep m) =>-  SegOpLevel rep ->-  Pat (LetDec rep) ->-  Certs ->-  SubExp -> -- segment size-  [SegBinOp rep] ->-  Lambda rep ->-  [VName] ->-  [(VName, SubExp)] -> -- ispace = pair of (gtid, size) for the maps on "top" of this reduction-  [KernelInput] -> -- inps = inputs that can be looked up by using the gtids from ispace-  m (Stms rep)-segRed lvl pat cs w ops map_lam arrs ispace inps = runBuilder_ $ do-  (kspace, kbody) <- prepareRedOrScan cs w map_lam arrs ispace inps-  letBind pat . Op . segOp $-    SegRed lvl kspace (lambdaReturnType map_lam) kbody ops--segScan ::-  (MonadFreshNames m, DistRep rep, HasScope rep m) =>-  SegOpLevel rep ->-  Pat (LetDec rep) ->-  Certs ->-  SubExp -> -- segment size-  [SegBinOp rep] ->-  Lambda rep ->-  SegPostOp rep ->-  [VName] ->-  [(VName, SubExp)] -> -- ispace = pair of (gtid, size) for the maps on "top" of this scan-  [KernelInput] -> -- inps = inputs that can be looked up by using the gtids from ispace-  m (Stms rep)-segScan lvl pat cs w ops map_lam post_op arrs ispace inps = runBuilder_ $ do-  let SegPostOp post_lam = post_op-  (kspace, kbody) <- prepareRedOrScan cs w map_lam arrs ispace inps-  post_lam' <- runLambdaBuilder (lambdaParams post_lam) $ do-    mapM_ readKernelInput inps-    bodyBind $ lambdaBody post_lam-  letBind pat . Op . segOp $-    SegScan lvl kspace (lambdaReturnType map_lam) kbody ops (SegPostOp post_lam')--segMap ::-  (MonadFreshNames m, DistRep rep, HasScope rep m) =>-  SegOpLevel rep ->-  Pat (LetDec rep) ->-  SubExp -> -- segment size-  Lambda rep ->-  [VName] ->-  [(VName, SubExp)] -> -- ispace = pair of (gtid, size) for the maps on "top" of this map-  [KernelInput] -> -- inps = inputs that can be looked up by using the gtids from ispace-  m (Stms rep)-segMap lvl pat w map_lam arrs ispace inps = runBuilder_ $ do-  (kspace, kbody) <- prepareRedOrScan mempty w map_lam arrs ispace inps-  letBind pat . Op . segOp $-    SegMap lvl kspace (lambdaReturnType map_lam) kbody--dummyDim ::-  (MonadBuilder m) =>-  Pat Type ->-  m (Pat Type, [(VName, SubExp)], m ())-dummyDim pat = do-  -- We add a unit-size segment on top to ensure that the result-  -- of the SegRed is an array, which we then immediately index.-  -- This is useful in the case that the value is used on the-  -- device afterwards, as this may save an expensive-  -- host-device copy (scalars are kept on the host, but arrays-  -- may be on the device).-  let addDummyDim t = t `arrayOfRow` intConst Int64 1-  pat' <- fmap addDummyDim <$> renamePat pat-  dummy <- newVName "dummy"-  let ispace = [(dummy, intConst Int64 1)]--  pure-    ( pat',-      ispace,-      forM_ (zip (patNames pat') (patNames pat)) $ \(from, to) -> do-        from_t <- lookupType from-        letBindNames [to] . BasicOp $-          case from_t of-            Acc {} -> SubExp $ Var from-            _ -> Index from $ fullSlice from_t [DimFix $ intConst Int64 0]-    )--nonSegRed ::-  (MonadFreshNames m, DistRep rep, HasScope rep m) =>-  SegOpLevel rep ->-  Pat Type ->-  SubExp ->-  [SegBinOp rep] ->-  Lambda rep ->-  [VName] ->-  m (Stms rep)-nonSegRed lvl pat w ops map_lam arrs = runBuilder_ $ do-  (pat', ispace, read_dummy) <- dummyDim pat-  addStms =<< segRed lvl pat' mempty w ops map_lam arrs ispace []-  read_dummy--segHist ::-  (DistRep rep, MonadFreshNames m, HasScope rep m) =>-  SegOpLevel rep ->-  Pat Type ->-  SubExp ->-  -- | Segment indexes and sizes.-  [(VName, SubExp)] ->-  [KernelInput] ->-  [HistOp rep] ->-  Lambda rep ->-  [VName] ->-  m (Stms rep)-segHist lvl pat arr_w ispace inps ops lam arrs = runBuilder_ $ do-  gtid <- newVName "gtid"-  space <- mkSegSpace $ ispace ++ [(gtid, arr_w)]--  kbody <- fmap (uncurry (flip $ Body ())) $-    runBuilder $-      localScope (scopeOfSegSpace space) $ do-        mapM_ readKernelInput inps-        forM_ (zip (lambdaParams lam) arrs) $ \(p, arr) -> do-          arr_t <- lookupType arr-          letBindNames [paramName p] $-            BasicOp $-              Index arr $-                fullSlice arr_t [DimFix $ Var gtid]-        res <- bodyBind (lambdaBody lam)-        forM res $ \(SubExpRes cs se) ->-          pure $ Returns ResultMaySimplify cs se--  letBind pat $ Op $ segOp $ SegHist lvl space (lambdaReturnType lam) kbody ops--mapKernelSkeleton ::-  (DistRep rep, HasScope rep m, MonadFreshNames m) =>-  [(VName, SubExp)] ->-  [KernelInput] ->-  m (SegSpace, Stms rep)-mapKernelSkeleton ispace inputs = do-  read_input_stms <- runBuilder_ $ mapM readKernelInput inputs--  space <- mkSegSpace ispace-  pure (space, read_input_stms)--mapKernel ::-  (DistRep rep, HasScope rep m, MonadFreshNames m) =>-  MkSegLevel rep m ->-  [(VName, SubExp)] ->-  [KernelInput] ->-  [Type] ->-  KernelBody rep ->-  m (SegOp (SegOpLevel rep) rep, Stms rep)-mapKernel mk_lvl ispace inputs rts (Body () kstms krets) = runBuilderT' $ do-  (space, read_input_stms) <- mapKernelSkeleton ispace inputs--  let kbody' = Body () (read_input_stms <> kstms) krets--  -- If the kernel creates arrays (meaning it will require memory-  -- expansion), we want to truncate the amount of threads.-  -- Otherwise, have at it!  This is a bit of a hack - in principle,-  -- we should make this decision later, when we have a clearer idea-  -- of what is happening inside the kernel.-  let r = if all primType rts then ManyThreads else NoRecommendation SegVirt--  lvl <- mk_lvl (map snd ispace) "segmap" r--  pure $ SegMap lvl space rts kbody'--data KernelInput = KernelInput-  { kernelInputName :: VName,-    kernelInputType :: Type,-    kernelInputArray :: VName,-    kernelInputIndices :: [SubExp]-  }-  deriving (Show)--readKernelInput ::-  (DistRep (Rep m), MonadBuilder m) =>-  KernelInput ->-  m ()-readKernelInput inp = do-  let pe = PatElem (kernelInputName inp) $ kernelInputType inp-  letBind (Pat [pe]) . BasicOp $-    case kernelInputType inp of-      Acc {} ->-        SubExp $ Var $ kernelInputArray inp-      _ ->-        Index (kernelInputArray inp) . Slice $-          map DimFix (kernelInputIndices inp)-            ++ map sliceDim (arrayDims (kernelInputType inp))
− src/Futhark/Pass/ExtractKernels/DistributeNests.hs 
@@ -1,1147 +0,0 @@-{-# LANGUAGE LambdaCase #-}-{-# LANGUAGE TypeFamilies #-}-{-# OPTIONS_GHC -Wno-overlapping-patterns -Wno-incomplete-patterns -Wno-incomplete-uni-patterns -Wno-incomplete-record-updates #-}--module Futhark.Pass.ExtractKernels.DistributeNests-  ( MapLoop (..),-    mapLoopStm,-    bodyContainsParallelism,-    lambdaContainsParallelism,-    determineReduceOp,-    histKernel,-    DistEnv (..),-    DistAcc (..),-    runDistNestT,-    DistNestT,-    liftInner,-    distributeMap,-    distribute,-    distributeSingleStm,-    distributeMapBodyStms,-    addStmsToAcc,-    addStmToAcc,-    permutationAndMissing,-    addPostStms,-    postStm,-    inNesting,-  )-where--import Control.Arrow (first)-import Control.Monad-import Control.Monad.RWS.Strict-import Control.Monad.Reader-import Control.Monad.Trans.Maybe-import Control.Monad.Writer.Strict-import Data.List (find, partition, tails)-import Data.List.NonEmpty (NonEmpty (..))-import Data.Map qualified as M-import Data.Maybe-import Futhark.IR-import Futhark.IR.GPU.Op (SegVirt (..))-import Futhark.IR.SOACS (SOACS)-import Futhark.IR.SOACS qualified as SOACS-import Futhark.IR.SOACS.SOAC hiding (HistOp, histDest)-import Futhark.IR.SOACS.Simplify (simpleSOACS, simplifyStms)-import Futhark.IR.SegOp-import Futhark.MonadFreshNames-import Futhark.Pass.ExtractKernels.BlockedKernel-import Futhark.Pass.ExtractKernels.Distribution-import Futhark.Pass.ExtractKernels.ISRWIM-import Futhark.Pass.ExtractKernels.Interchange-import Futhark.Tools-import Futhark.Transform.CopyPropagate-import Futhark.Transform.FirstOrderTransform qualified as FOT-import Futhark.Transform.Rename-import Futhark.Util.Log--scopeForSOACs :: (SameScope rep SOACS) => Scope rep -> Scope SOACS-scopeForSOACs = castScope--data MapLoop = MapLoop (Pat Type) (StmAux ()) SubExp (Lambda SOACS) [VName]--mapLoopStm :: (MonadFreshNames m) => MapLoop -> m (Stm SOACS)-mapLoopStm (MapLoop pat aux w lam arrs) =-  Let pat aux . Op . Screma w arrs <$> mapSOAC lam--data DistEnv rep m = DistEnv-  { distNest :: Nestings,-    distScope :: Scope rep,-    distOnTopLevelStms :: Stms SOACS -> DistNestT rep m (Stms rep),-    distOnInnerMap ::-      MapLoop ->-      DistAcc rep ->-      DistNestT rep m (DistAcc rep),-    distOnSOACSStms :: Stm SOACS -> Builder rep (Stms rep),-    distOnSOACSLambda :: Lambda SOACS -> Builder rep (Lambda rep),-    distSegLevel :: MkSegLevel rep m-  }--data DistAcc rep = DistAcc-  { distTargets :: Targets,-    distStms :: Stms rep-  }--data DistRes rep = DistRes-  { accPostStms :: PostStms rep,-    accLog :: Log-  }--instance Semigroup (DistRes rep) where-  DistRes ks1 log1 <> DistRes ks2 log2 =-    DistRes (ks1 <> ks2) (log1 <> log2)--instance Monoid (DistRes rep) where-  mempty = DistRes mempty mempty--newtype PostStms rep = PostStms {unPostStms :: Stms rep}--instance Semigroup (PostStms rep) where-  PostStms xs <> PostStms ys = PostStms $ ys <> xs--instance Monoid (PostStms rep) where-  mempty = PostStms mempty--typeEnvFromDistAcc :: (DistRep rep) => DistAcc rep -> Scope rep-typeEnvFromDistAcc = scopeOfPat . fst . outerTarget . distTargets--addStmsToAcc :: Stms rep -> DistAcc rep -> DistAcc rep-addStmsToAcc stms acc =-  acc {distStms = stms <> distStms acc}--addStmToAcc ::-  (MonadFreshNames m, DistRep rep) =>-  Stm SOACS ->-  DistAcc rep ->-  DistNestT rep m (DistAcc rep)-addStmToAcc stm acc = do-  onSoacs <- asks distOnSOACSStms-  (stm', _) <- runBuilder $ onSoacs stm-  pure acc {distStms = stm' <> distStms acc}--soacsLambda ::-  (MonadFreshNames m, DistRep rep) =>-  Lambda SOACS ->-  DistNestT rep m (Lambda rep)-soacsLambda lam = do-  onLambda <- asks distOnSOACSLambda-  fst <$> runBuilder (onLambda lam)--newtype DistNestT rep m a-  = DistNestT (ReaderT (DistEnv rep m) (WriterT (DistRes rep) m) a)-  deriving-    ( Functor,-      Applicative,-      Monad,-      MonadReader (DistEnv rep m),-      MonadWriter (DistRes rep)-    )--liftInner :: (LocalScope rep m, DistRep rep) => m a -> DistNestT rep m a-liftInner m = do-  outer_scope <- askScope-  DistNestT $-    lift $-      lift $ do-        inner_scope <- askScope-        localScope (outer_scope `M.difference` inner_scope) m--instance (MonadFreshNames m) => MonadFreshNames (DistNestT rep m) where-  getNameSource = DistNestT $ lift getNameSource-  putNameSource = DistNestT . lift . putNameSource--instance (Monad m, ASTRep rep) => HasScope rep (DistNestT rep m) where-  askScope = asks distScope--instance (Monad m, ASTRep rep) => LocalScope rep (DistNestT rep m) where-  localScope types = local $ \env ->-    env {distScope = types <> distScope env}--instance (Monad m) => MonadLogger (DistNestT rep m) where-  addLog msgs = tell mempty {accLog = msgs}--runDistNestT ::-  (MonadLogger m, DistRep rep) =>-  DistEnv rep m ->-  DistNestT rep m (DistAcc rep) ->-  m (Stms rep)-runDistNestT env (DistNestT m) = do-  (acc, res) <- runWriterT $ runReaderT m env-  addLog $ accLog res-  -- There may be a few final targets remaining - these correspond to-  -- arrays that are identity mapped, and must have statements-  -- inserted here.-  pure $-    unPostStms (accPostStms res) <> identityStms (outerTarget $ distTargets acc)-  where-    outermost = nestingLoop $-      case distNest env of-        (nest, []) -> nest-        (_, nest : _) -> nest-    params_to_arrs =-      map (first paramName) $-        loopNestingParamsAndArrs outermost--    identityStms (rem_pat, res) =-      stmsFromList $ zipWith identityStm (patElems rem_pat) res-    identityStm pe (SubExpRes cs (Var v))-      | Just arr <- lookup v params_to_arrs =-          certify cs . Let (Pat [pe]) (defAux ()) . BasicOp $-            Replicate mempty (Var arr)-    identityStm pe (SubExpRes cs se) =-      certify cs . Let (Pat [pe]) (defAux ()) . BasicOp $-        Replicate (Shape [loopNestingWidth outermost]) se--addPostStms :: (Monad m) => PostStms rep -> DistNestT rep m ()-addPostStms ks = tell $ mempty {accPostStms = ks}--postStm :: (Monad m) => Stms rep -> DistNestT rep m ()-postStm stms = addPostStms $ PostStms stms--withStm ::-  (Monad m, DistRep rep) =>-  Stm SOACS ->-  DistNestT rep m a ->-  DistNestT rep m a-withStm stm = local $ \env ->-  env-    { distScope =-        castScope (scopeOf stm) <> distScope env,-      distNest =-        letBindInInnerNesting provided $-          distNest env-    }-  where-    provided = namesFromList $ patNames $ stmPat stm--leavingNesting ::-  (MonadFreshNames m, DistRep rep) =>-  DistAcc rep ->-  DistNestT rep m (DistAcc rep)-leavingNesting acc =-  case popInnerTarget $ distTargets acc of-    Nothing ->-      error "The kernel targets list is unexpectedly small"-    Just ((pat, res), newtargets)-      | not $ null $ distStms acc -> do-          -- Any statements left over correspond to something that-          -- could not be distributed because it would cause irregular-          -- arrays.  These must be reconstructed into a a Map SOAC-          -- that will be sequentialised. XXX: life would be better if-          -- we were able to distribute irregular parallelism.-          (Nesting _ inner, _) <- asks distNest-          let MapNesting _ aux w params_and_arrs = inner-              body = Body () (distStms acc) res-              used_in_body = freeIn body-              (used_params, used_arrs) =-                unzip $-                  filter ((`nameIn` used_in_body) . paramName . fst) params_and_arrs-              lam' =-                Lambda-                  { lambdaParams = used_params,-                    lambdaBody = body,-                    lambdaReturnType = map rowType $ patTypes pat-                  }-          stms <--            runBuilder_-              . auxing aux-              . FOT.transformSOAC pat-              . Screma w used_arrs-              =<< mapSOAC lam'--          pure $ acc {distTargets = newtargets, distStms = stms}-      | otherwise -> do-          -- Any results left over correspond to a Replicate or a Copy in-          -- the parent nesting, depending on whether the argument is a-          -- parameter of the innermost nesting.-          (Nesting _ inner_nesting, _) <- asks distNest-          let w = loopNestingWidth inner_nesting-              aux = loopNestingAux inner_nesting-              inps = loopNestingParamsAndArrs inner_nesting--              remnantStm pe (SubExpRes cs (Var v))-                | Just (_, arr) <- find ((== v) . paramName . fst) inps =-                    certify cs . Let (Pat [pe]) aux . BasicOp $-                      Replicate mempty (Var arr)-              remnantStm pe (SubExpRes cs se) =-                certify cs . Let (Pat [pe]) aux . BasicOp $-                  Replicate (Shape [w]) se--              stms =-                stmsFromList $ zipWith remnantStm (patElems pat) res--          pure $ acc {distTargets = newtargets, distStms = stms}--mapNesting ::-  (MonadFreshNames m, DistRep rep) =>-  Pat Type ->-  StmAux () ->-  SubExp ->-  Lambda SOACS ->-  [VName] ->-  DistNestT rep m (DistAcc rep) ->-  DistNestT rep m (DistAcc rep)-mapNesting pat aux w lam arrs m =-  local extend $ leavingNesting =<< m-  where-    nest =-      Nesting mempty $-        MapNesting pat aux w $-          zip (lambdaParams lam) arrs-    extend env =-      env-        { distNest = pushInnerNesting nest $ distNest env,-          distScope = castScope (scopeOf lam) <> distScope env-        }--inNesting ::-  (Monad m, DistRep rep) =>-  KernelNest ->-  DistNestT rep m a ->-  DistNestT rep m a-inNesting (outer, nests) = local $ \env ->-  env-    { distNest = (inner, nests'),-      distScope = foldMap scopeOfLoopNesting (outer : nests) <> distScope env-    }-  where-    (inner, nests') =-      case reverse nests of-        [] -> (asNesting outer, [])-        (inner' : ns) -> (asNesting inner', map asNesting $ outer : reverse ns)-    asNesting = Nesting mempty--bodyContainsParallelism :: Body SOACS -> Bool-bodyContainsParallelism = any isParallelStm . bodyStms-  where-    isParallelStm stm =-      isMap (stmExp stm)-        && not ("sequential" `inAttrs` stmAuxAttrs (stmAux stm))-    isMap BasicOp {} = False-    isMap Apply {} = False-    isMap Match {} = False-    isMap (Loop _ ForLoop {} body) = bodyContainsParallelism body-    isMap (Loop _ WhileLoop {} _) = False-    isMap (WithAcc _ lam) = bodyContainsParallelism $ lambdaBody lam-    isMap Op {} = True--lambdaContainsParallelism :: Lambda SOACS -> Bool-lambdaContainsParallelism = bodyContainsParallelism . lambdaBody--distributeMapBodyStms ::-  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>-  DistAcc rep ->-  Stms SOACS ->-  DistNestT rep m (DistAcc rep)-distributeMapBodyStms orig_acc = distribute <=< onStms orig_acc . stmsToList-  where-    onStms acc [] = pure acc-    onStms acc (Let pat aux (Op (Stream w arrs accs lam)) : stms) = do-      types <- asksScope scopeForSOACs-      stream_stms <--        snd <$> runBuilderT (sequentialStreamWholeArray pat w accs lam arrs) types-      stream_stms' <--        runReaderT (copyPropagateInStms simpleSOACS types stream_stms) types-      onStms acc $ stmsToList (fmap (certify (stmAuxCerts aux)) stream_stms') ++ stms-    onStms acc (stm : stms) =-      -- It is important that stm is in scope if 'maybeDistributeStm'-      -- wants to distribute, even if this causes the slightly silly-      -- situation that stm is in scope of itself.-      withStm stm $ maybeDistributeStm stm =<< onStms acc stms--onInnerMap :: (Monad m) => MapLoop -> DistAcc rep -> DistNestT rep m (DistAcc rep)-onInnerMap loop acc = do-  f <- asks distOnInnerMap-  f loop acc--onTopLevelStms :: (Monad m) => Stms SOACS -> DistNestT rep m ()-onTopLevelStms stms = do-  f <- asks distOnTopLevelStms-  postStm =<< f stms--maybeDistributeStm ::-  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>-  Stm SOACS ->-  DistAcc rep ->-  DistNestT rep m (DistAcc rep)-maybeDistributeStm stm acc-  | "sequential" `inAttrs` stmAuxAttrs (stmAux stm) =-      addStmToAcc stm acc-maybeDistributeStm (Let pat aux (Op soac)) acc-  | "sequential_outer" `inAttrs` stmAuxAttrs aux =-      distributeMapBodyStms acc . fmap (certify (stmAuxCerts aux))-        =<< runBuilder_ (FOT.transformSOAC pat soac)-maybeDistributeStm stm@(Let pat _ (Op (Screma w arrs form))) acc-  | Just lam <- isMapSOAC form =-      -- Only distribute inside the map if we can distribute everything-      -- following the map.-      distributeIfPossible acc >>= \case-        Nothing -> addStmToAcc stm acc-        Just acc' -> distribute =<< onInnerMap (MapLoop pat (stmAux stm) w lam arrs) acc'-maybeDistributeStm stm@(Let pat aux (Loop merge form@ForLoop {} body)) acc-  | all (`notNameIn` freeIn (patTypes pat)) (patNames pat),-    bodyContainsParallelism body =-      distributeSingleStm acc stm >>= \case-        Just (kernels, res, nest, acc')-          | -- XXX: We cannot distribute if this loop depends on-            -- certificates bound within the loop nest (well, we could,-            -- but interchange would not be valid).  This is not a-            -- fundamental restriction, but an artifact of our-            -- certificate representation, which we should probably-            -- rethink.-            not $-              (freeIn form <> freeIn aux)-                `namesIntersect` boundInKernelNest nest,-            Just (perm, pat_unused) <- permutationAndMissing pat res ->-              -- We need to pretend pat_unused was used anyway, by adding-              -- it to the kernel nest.-              localScope (typeEnvFromDistAcc acc') $ do-                addPostStms kernels-                nest' <- expandKernelNest pat_unused nest-                types <- asksScope scopeForSOACs--                -- Simplification is key to hoisting out statements that-                -- were variant to the loop, but invariant to the outer maps-                -- (which are now innermost).-                stms <--                  (`runReaderT` types) $-                    simplifyStms =<< interchangeLoops nest' (SeqLoop perm pat merge form body)-                onTopLevelStms stms-                pure acc'-        _ ->-          addStmToAcc stm acc-maybeDistributeStm stm@(Let pat _ (Match cond cases defbody ret)) acc-  | all (`notNameIn` freeIn pat) (patNames pat),-    any bodyContainsParallelism (defbody : map caseBody cases)-      || not (all primType (matchReturns ret)) =-      distributeSingleStm acc stm >>= \case-        Just (kernels, res, nest, acc')-          | not $-              (freeIn cond <> freeIn ret) `namesIntersect` boundInKernelNest nest,-            Just (perm, pat_unused) <- permutationAndMissing pat res ->-              -- We need to pretend pat_unused was used anyway, by adding-              -- it to the kernel nest.-              localScope (typeEnvFromDistAcc acc') $ do-                nest' <- expandKernelNest pat_unused nest-                addPostStms kernels-                types <- asksScope scopeForSOACs-                let branch = Branch perm pat cond cases defbody ret-                stms <--                  (`runReaderT` types) $-                    simplifyStms . oneStm =<< interchangeBranch nest' branch-                onTopLevelStms stms-                pure acc'-        _ ->-          addStmToAcc stm acc-maybeDistributeStm stm@(Let pat _ (WithAcc inputs lam)) acc-  | lambdaContainsParallelism lam =-      distributeSingleStm acc stm >>= \case-        Just (kernels, res, nest, acc')-          | not $-              freeIn (drop num_accs (lambdaReturnType lam))-                `namesIntersect` boundInKernelNest nest,-            Just (perm, pat_unused) <- permutationAndMissing pat res ->-              -- We need to pretend pat_unused was used anyway, by adding-              -- it to the kernel nest.-              localScope (typeEnvFromDistAcc acc') $ do-                nest' <- expandKernelNest pat_unused nest-                types <- asksScope scopeForSOACs-                addPostStms kernels-                let withacc = WithAccStm perm pat inputs lam-                stms <--                  (`runReaderT` types) $-                    simplifyStms . oneStm =<< interchangeWithAcc nest' withacc-                onTopLevelStms stms-                pure acc'-        _ ->-          addStmToAcc stm acc-  where-    num_accs = length inputs-maybeDistributeStm (Let pat aux (Op (Screma w arrs form))) acc-  | Just [Reduce comm lam nes] <- isReduceSOAC form,-    Just m <- irwim pat w comm lam $ zip nes arrs = do-      types <- asksScope scopeForSOACs-      (_, stms) <- runBuilderT (auxing aux m) types-      distributeMapBodyStms acc stms---- Parallelise segmented Hist.-maybeDistributeStm stm@(Let pat aux (Op (Hist w as ops lam))) acc =-  distributeSingleStm acc stm >>= \case-    Just (kernels, res, nest, acc')-      | Just (perm, pat_unused) <- permutationAndMissing pat res ->-          localScope (typeEnvFromDistAcc acc') $ do-            lam' <- soacsLambda lam-            nest' <- expandKernelNest pat_unused nest-            addPostStms kernels-            postStm =<< segmentedHistKernel nest' perm (stmAuxCerts aux) w ops lam' as-            pure acc'-    _ ->-      addStmToAcc stm acc--- Parallelise Index slices if the result is going to be returned--- directly from the kernel.  This is because we would otherwise have--- to sequentialise writing the result, which may be costly.-maybeDistributeStm stm@(Let (Pat [pe]) aux (BasicOp (Index arr slice))) acc-  | not $ null $ sliceDims slice,-    Var (patElemName pe) `elem` map resSubExp (snd (innerTarget (distTargets acc))) =-      distributeSingleStm acc stm >>= \case-        Just (kernels, _res, nest, acc') ->-          localScope (typeEnvFromDistAcc acc') $ do-            addPostStms kernels-            postStm =<< segmentedGatherKernel nest (stmAuxCerts aux) arr slice-            pure acc'-        _ ->-          addStmToAcc stm acc--- If the scan can be distributed by itself, we will turn it into a--- segmented scan.------ If the scan cannot be distributed by itself, it will be--- sequentialised in the default case for this function.-maybeDistributeStm stm@(Let pat aux (Op (Screma w arrs form))) acc-  | Just (post_lam, scans, map_lam) <- isMaposcanomapSOAC form,-    Scan op_lam nes <- singleScan scans =-      distributeSingleStm acc stm >>= \case-        Just (kernels, res, nest, acc')-          | Just (perm, pat_unused) <- permutationAndMissing pat res ->-              -- We need to pretend pat_unused was used anyway, by adding-              -- it to the kernel nest.-              localScope (typeEnvFromDistAcc acc') $ do-                nest' <- expandKernelNest pat_unused nest-                map_lam' <- soacsLambda map_lam-                localScope (typeEnvFromDistAcc acc') $-                  segmentedScanomapKernel nest' perm (stmAuxCerts aux) w op_lam map_lam' post_lam nes arrs-                    >>= kernelOrNot mempty stm acc kernels acc'-        _ ->-          addStmToAcc stm acc--- If the map function of the reduction contains parallelism we split--- it, so that the parallelism can be exploited.-maybeDistributeStm (Let pat aux (Op (Screma w arrs form))) acc-  | Just (reds, map_lam) <- isRedomapSOAC form,-    lambdaContainsParallelism map_lam = do-      (mapstm, redstm) <--        redomapToMapAndReduce pat (w, reds, map_lam, arrs)-      distributeMapBodyStms acc $ oneStm mapstm {stmAux = aux} <> oneStm redstm--- if the reduction can be distributed by itself, we will turn it into a--- segmented reduce.------ If the reduction cannot be distributed by itself, it will be--- sequentialised in the default case for this function.-maybeDistributeStm stm@(Let pat aux (Op (Screma w arrs form))) acc-  | Just (reds, map_lam) <- isRedomapSOAC form,-    Reduce comm lam nes <- singleReduce reds =-      distributeSingleStm acc stm >>= \case-        Just (kernels, res, nest, acc')-          | Just (perm, pat_unused) <- permutationAndMissing pat res ->-              -- We need to pretend pat_unused was used anyway, by adding-              -- it to the kernel nest.-              localScope (typeEnvFromDistAcc acc') $ do-                nest' <- expandKernelNest pat_unused nest--                lam' <- soacsLambda lam-                map_lam' <- soacsLambda map_lam--                let comm'-                      | commutativeLambda lam = Commutative-                      | otherwise = comm--                regularSegmentedRedomapKernel nest' perm (stmAuxCerts aux) w comm' lam' map_lam' nes arrs-                  >>= kernelOrNot mempty stm acc kernels acc'-        _ ->-          addStmToAcc stm acc-maybeDistributeStm (Let pat aux (Op (Screma w arrs form))) acc = do-  -- This Screma is too complicated for us to immediately do-  -- anything, so split it up and try again.-  scope <- asksScope scopeForSOACs-  distributeMapBodyStms acc . fmap (certify (stmAuxCerts aux)) . snd-    =<< runBuilderT (dissectScrema pat w form arrs) scope-maybeDistributeStm stm@(Let _ aux (BasicOp (Replicate shape (Var stm_arr)))) acc = do-  distributeSingleUnaryStm acc stm stm_arr $ \nest outerpat arr ->-    if shape == mempty-      then pure $ oneStm $ Let outerpat aux $ BasicOp $ Replicate mempty $ Var arr-      else runBuilder_ $ auxing aux $ do-        arr_t <- lookupType arr-        let arr_r = arrayRank arr_t-            nest_r = length (snd nest) + 1-            res_r = arr_r + shapeRank shape-        -- Move the to-be-replicated dimensions outermost.-        arr_tr <--          letExp (baseName arr <> "_tr") . BasicOp $-            Rearrange arr ([nest_r .. arr_r - 1] ++ [0 .. nest_r - 1])-        -- Replicate the now-outermost dimensions appropriately.-        arr_tr_rep <--          letExp (baseName arr <> "_tr_rep") . BasicOp $-            Replicate shape (Var arr_tr)-        -- Move the replicated dimensions back where they belong.-        letBind outerpat . BasicOp $-          Rearrange arr_tr_rep ([res_r - nest_r .. res_r - 1] ++ [0 .. res_r - nest_r - 1])-maybeDistributeStm stm@(Let _ aux (BasicOp (Replicate shape v))) acc = do-  distributeSingleStm acc stm >>= \case-    Just (kernels, _, nest, acc')-      | boundInKernelNest nest == mempty -> do-          addPostStms kernels-          let outerpat = loopNestingPat $ fst nest-              nest_shape = Shape $ kernelNestWidths nest-          localScope (typeEnvFromDistAcc acc') $ do-            postStm <=< runBuilder_ . auxing aux . letBind outerpat $-              BasicOp (Replicate (nest_shape <> shape) v)-            pure acc'-    _ -> addStmToAcc stm acc--- Opaques are applied to the full array, because otherwise they can--- drastically inhibit parallelisation in some cases.-maybeDistributeStm stm@(Let (Pat [pe]) aux (BasicOp (Opaque _ (Var stm_arr)))) acc-  | not $ primType $ typeOf pe =-      distributeSingleUnaryStm acc stm stm_arr $ \_ outerpat arr ->-        pure $ oneStm $ Let outerpat aux $ BasicOp $ Replicate mempty $ Var arr-maybeDistributeStm stm@(Let _ aux (BasicOp (Rearrange stm_arr perm))) acc =-  distributeSingleUnaryStm acc stm stm_arr $ \nest outerpat arr -> do-    let r = length (snd nest) + 1-        perm' = [0 .. r - 1] ++ map (+ r) perm-    -- We need to add a copy, because the original map nest-    -- will have produced an array without aliases, and so must we.-    arr' <- newName arr-    arr_t <- lookupType arr-    pure $-      stmsFromList-        [ Let (Pat [PatElem arr' arr_t]) aux $ BasicOp $ Replicate mempty $ Var arr,-          Let outerpat aux $ BasicOp $ Rearrange arr' perm'-        ]-maybeDistributeStm stm@(Let _ aux (BasicOp (Reshape stm_arr reshape))) acc =-  distributeSingleUnaryStm acc stm stm_arr $ \nest outerpat arr -> do-    let outer = Shape (kernelNestWidths nest)-        reshape' = reshapeCoerce outer <> newshapeInner outer reshape-    pure $ oneStm $ Let outerpat aux $ BasicOp $ Reshape arr reshape'-maybeDistributeStm stm@(Let pat aux (BasicOp (Update _ arr slice (Var v)))) acc-  | not $ null $ sliceDims slice =-      distributeSingleStm acc stm >>= \case-        Just (kernels, res, nest, acc')-          | map resSubExp res == map Var (patNames $ stmPat stm),-            Just (perm, pat_unused) <- permutationAndMissing pat res -> do-              addPostStms kernels-              localScope (typeEnvFromDistAcc acc') $ do-                nest' <- expandKernelNest pat_unused nest-                postStm-                  =<< segmentedUpdateKernel nest' perm (stmAuxCerts aux) arr slice v-                pure acc'-        _ -> addStmToAcc stm acc-maybeDistributeStm stm@(Let _ _ (BasicOp (Concat d (x :| xs) w))) acc =-  distributeSingleStm acc stm >>= \case-    Just (kernels, _, nest, acc') ->-      localScope (typeEnvFromDistAcc acc') $-        segmentedConcat nest-          >>= kernelOrNot mempty stm acc kernels acc'-    _ ->-      addStmToAcc stm acc-  where-    segmentedConcat nest =-      isSegmentedOp nest [0] mempty mempty [] (x : xs) $-        \pat _ _ _ (x' : xs') ->-          let d' = d + length (snd nest) + 1-           in addStm $ Let pat mempty $ BasicOp $ Concat d' (x' :| xs') w-maybeDistributeStm stm acc =-  addStmToAcc stm acc--distributeSingleUnaryStm ::-  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>-  DistAcc rep ->-  Stm SOACS ->-  VName ->-  (KernelNest -> Pat Type -> VName -> DistNestT rep m (Stms rep)) ->-  DistNestT rep m (DistAcc rep)-distributeSingleUnaryStm acc stm stm_arr f =-  distributeSingleStm acc stm >>= \case-    Just (kernels, res, nest, acc')-      | map resSubExp res == map Var (patNames $ stmPat stm),-        (outer, _) <- nest,-        [(_, arr)] <- loopNestingParamsAndArrs outer,-        boundInKernelNest nest `namesIntersection` freeIn stm-          == oneName stm_arr,-        perfectlyMapped arr nest -> do-          addPostStms kernels-          let outerpat = loopNestingPat $ fst nest-          localScope (typeEnvFromDistAcc acc') $ do-            postStm =<< f nest outerpat arr-            pure acc'-    _ -> addStmToAcc stm acc-  where-    perfectlyMapped arr (outer, nest)-      | [(p, arr')] <- loopNestingParamsAndArrs outer,-        arr == arr' =-          case nest of-            [] -> paramName p == stm_arr-            x : xs -> perfectlyMapped (paramName p) (x, xs)-      | otherwise =-          False--distribute ::-  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>-  DistAcc rep ->-  DistNestT rep m (DistAcc rep)-distribute acc =-  fromMaybe acc <$> distributeIfPossible acc--mkSegLevel ::-  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>-  DistNestT rep m (MkSegLevel rep (DistNestT rep m))-mkSegLevel = do-  mk_lvl <- asks distSegLevel-  pure $ \w desc r -> do-    (lvl, stms) <- lift $ liftInner $ runBuilderT' $ mk_lvl w desc r-    addStms stms-    pure lvl--distributeIfPossible ::-  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>-  DistAcc rep ->-  DistNestT rep m (Maybe (DistAcc rep))-distributeIfPossible acc = do-  nest <- asks distNest-  mk_lvl <- mkSegLevel-  tryDistribute mk_lvl nest (distTargets acc) (distStms acc) >>= \case-    Nothing -> pure Nothing-    Just (targets, kernel) -> do-      postStm kernel-      pure $-        Just-          DistAcc-            { distTargets = targets,-              distStms = mempty-            }--distributeSingleStm ::-  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>-  DistAcc rep ->-  Stm SOACS ->-  DistNestT-    rep-    m-    ( Maybe-        ( PostStms rep,-          Result,-          KernelNest,-          DistAcc rep-        )-    )-distributeSingleStm acc stm = do-  nest <- asks distNest-  mk_lvl <- mkSegLevel-  tryDistribute mk_lvl nest (distTargets acc) (distStms acc) >>= \case-    Nothing -> pure Nothing-    Just (targets, distributed_stms) ->-      tryDistributeStm nest targets stm >>= \case-        Nothing -> pure Nothing-        Just (res, targets', new_kernel_nest) ->-          pure $-            Just-              ( PostStms distributed_stms,-                res,-                new_kernel_nest,-                DistAcc-                  { distTargets = targets',-                    distStms = mempty-                  }-              )--segmentedUpdateKernel ::-  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>-  KernelNest ->-  [Int] ->-  Certs ->-  VName ->-  Slice SubExp ->-  VName ->-  DistNestT rep m (Stms rep)-segmentedUpdateKernel nest perm cs arr slice v = runBuilderT'_ $ do-  (base_ispace, kernel_inps) <- flatKernel nest-  let rank = length base_ispace + length (sliceDims slice)-      arr' =-        maybe (error "incorrectly typed Update") kernelInputArray $-          find ((== arr) . kernelInputName) kernel_inps--  arr_t <- lookupType arr'-  let arr_rank = arrayRank arr_t-      remnant_dims = drop rank $ arrayDims arr_t--  e <- withAcc [arr'] arr_rank $ \ ~[acc] -> do-    let slice_dims = sliceDims slice-    slice_gtids <- replicateM (length slice_dims) (newVName "gtid_slice")-    remnant_gtids <- replicateM (length remnant_dims) $ newVName "gtid_remnant"--    let ispace =-          base_ispace-            <> zip slice_gtids slice_dims-            <> zip remnant_gtids remnant_dims--    body <- runBodyBuilder $ do-      -- Compute indexes into full array.-      v' <--        certifying cs . letSubExp "v" . BasicOp . Index v $-          Slice (map (DimFix . Var) slice_gtids)-      slice_is <--        traverse (toSubExp "index") $-          fixSlice (fmap pe64 slice) $-            map (pe64 . Var) slice_gtids--      let write_is = map (Var . fst) base_ispace ++ slice_is-      acc' <- letExp "acc" $ BasicOp $ UpdateAcc Safe acc write_is [v']-      pure [Returns ResultMaySimplify mempty $ Var acc']--    -- Remove unused kernel inputs, since some of these might-    -- reference the array we are scattering into.-    let kernel_inps' =-          filter ((`nameIn` freeIn body) . kernelInputName) kernel_inps--    mk_lvl <- lift mkSegLevel-    acc_t <- lookupType acc-    (k, prestms) <--      lift $ mapKernel mk_lvl ispace kernel_inps' [acc_t] body-    addStms prestms-    fmap pure $ letSubExp "segmented_upd" $ Op $ segOp k--  let pat = Pat . rearrangeShape perm $ patElems $ loopNestingPat $ fst nest-  letBind pat e--segmentedGatherKernel ::-  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>-  KernelNest ->-  Certs ->-  VName ->-  Slice SubExp ->-  DistNestT rep m (Stms rep)-segmentedGatherKernel nest cs arr slice = do-  let slice_dims = sliceDims slice-  slice_gtids <- replicateM (length slice_dims) (newVName "gtid_slice")--  (base_ispace, kernel_inps) <- flatKernel nest-  let ispace = base_ispace ++ zip slice_gtids slice_dims--  ((res_t, res), kstms) <- runBuilder $ do-    -- Compute indexes into full array.-    slice'' <--      subExpSlice . sliceSlice (primExpSlice slice) $-        primExpSlice $-          Slice $-            map (DimFix . Var) slice_gtids-    v' <- certifying cs $ letSubExp "v" $ BasicOp $ Index arr slice''-    v_t <- subExpType v'-    pure (v_t, Returns ResultMaySimplify mempty v')--  mk_lvl <- mkSegLevel-  (k, prestms) <--    mapKernel mk_lvl ispace kernel_inps [res_t] $-      Body () kstms [res]--  traverse renameStm <=< runBuilder_ $ do-    addStms prestms--    let pat = Pat $ patElems $ loopNestingPat $ fst nest--    letBind pat $ Op $ segOp k--segmentedHistKernel ::-  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>-  KernelNest ->-  [Int] ->-  Certs ->-  SubExp ->-  [SOACS.HistOp SOACS] ->-  Lambda rep ->-  [VName] ->-  DistNestT rep m (Stms rep)-segmentedHistKernel nest perm cs hist_w ops lam arrs = do-  -- We replicate some of the checking done by 'isSegmentedOp', but-  -- things are different because a Hist is not a reduction or-  -- scan.-  (ispace, inputs) <- flatKernel nest-  let orig_pat =-        Pat . rearrangeShape perm $-          patElems $-            loopNestingPat $-              fst nest--  -- The input/output arrays _must_ correspond to some kernel input,-  -- or else the original nested Hist would have been ill-typed.-  -- Find them.-  ops' <- forM ops $ \(SOACS.HistOp num_bins rf dests nes op) ->-    SOACS.HistOp num_bins rf-      <$> mapM (fmap kernelInputArray . findInput inputs) dests-      <*> pure nes-      <*> pure op--  mk_lvl <- asks distSegLevel-  onLambda <- asks distOnSOACSLambda-  let onLambda' = fmap fst . runBuilder . onLambda-  liftInner $-    runBuilderT'_ $ do-      -- It is important not to launch unnecessarily many threads for-      -- histograms, because it may mean we unnecessarily need to reduce-      -- subhistograms as well.-      lvl <- mk_lvl (hist_w : map snd ispace) "seghist" $ NoRecommendation SegNoVirt-      addStms-        =<< histKernel onLambda' lvl orig_pat ispace inputs cs hist_w ops' lam arrs-  where-    findInput kernel_inps a =-      maybe bad pure $ find ((== a) . kernelInputName) kernel_inps-    bad = error "Ill-typed nested Hist encountered."--histKernel ::-  (MonadBuilder m, DistRep (Rep m)) =>-  (Lambda SOACS -> m (Lambda (Rep m))) ->-  SegOpLevel (Rep m) ->-  Pat Type ->-  [(VName, SubExp)] ->-  [KernelInput] ->-  Certs ->-  SubExp ->-  [SOACS.HistOp SOACS] ->-  Lambda (Rep m) ->-  [VName] ->-  m (Stms (Rep m))-histKernel onLambda lvl orig_pat ispace inputs cs hist_w ops lam arrs = runBuilderT'_ $ do-  ops' <- forM ops $ \(SOACS.HistOp dest_shape rf dests nes op) -> do-    (op', nes', shape) <- determineReduceOp op nes-    op'' <- lift $ onLambda op'-    pure $ HistOp dest_shape rf dests nes' shape op''--  let isDest = flip elem $ concatMap histDest ops'-      inputs' = filter (not . isDest . kernelInputArray) inputs--  certifying cs $-    addStms-      =<< traverse renameStm-      =<< segHist lvl orig_pat hist_w ispace inputs' ops' lam arrs--determineReduceOp ::-  (MonadBuilder m) =>-  Lambda SOACS ->-  [SubExp] ->-  m (Lambda SOACS, [SubExp], Shape)-determineReduceOp lam nes =-  -- FIXME? We are assuming that the accumulator is a replicate, and-  -- we fish out its value in a gross way.-  case mapM subExpVar nes of-    Just ne_vs' -> do-      let (shape, lam') = isVectorMap lam-      nes' <- forM ne_vs' $ \ne_v -> do-        ne_v_t <- lookupType ne_v-        letSubExp "hist_ne" $-          BasicOp $-            Index ne_v $-              fullSlice ne_v_t $-                replicate (shapeRank shape) $-                  DimFix $-                    intConst Int64 0-      pure (lam', nes', shape)-    Nothing ->-      pure (lam, nes, mempty)--isVectorMap :: Lambda SOACS -> (Shape, Lambda SOACS)-isVectorMap lam-  | [Let (Pat pes) _ (Op (Screma w arrs form))] <--      stmsToList $ bodyStms $ lambdaBody lam,-    map resSubExp (bodyResult (lambdaBody lam)) == map (Var . patElemName) pes,-    Just map_lam <- isMapSOAC form,-    arrs == map paramName (lambdaParams lam) =-      let (shape, lam') = isVectorMap map_lam-       in (Shape [w] <> shape, lam')-  | otherwise = (mempty, lam)--segmentedScanomapKernel ::-  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>-  KernelNest ->-  [Int] ->-  Certs ->-  SubExp ->-  Lambda SOACS ->-  Lambda rep ->-  Lambda SOACS ->-  [SubExp] ->-  [VName] ->-  DistNestT rep m (Maybe (Stms rep))-segmentedScanomapKernel nest perm cs segment_size op_lam map_lam post_lam nes arrs = do-  mk_lvl <- asks distSegLevel-  onLambda <- asks distOnSOACSLambda-  let onLambda' = fmap fst . runBuilder . onLambda-  isSegmentedOp nest perm (freeIn op_lam) (freeIn map_lam) nes [] $-    \pat ispace inps nes' _ -> do-      (op_lam', nes'', shape) <- determineReduceOp op_lam nes'-      op_lam'' <- onLambda' op_lam'-      let scan_op = SegBinOp Noncommutative op_lam'' nes'' shape-      post_lam' <- onLambda' post_lam-      let post_op = SegPostOp post_lam'-      lvl <- mk_lvl (segment_size : map snd ispace) "segscan" $ NoRecommendation SegNoVirt-      addStms-        =<< traverse renameStm-        =<< segScan lvl pat cs segment_size [scan_op] map_lam post_op arrs ispace inps--regularSegmentedRedomapKernel ::-  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>-  KernelNest ->-  [Int] ->-  Certs ->-  SubExp ->-  Commutativity ->-  Lambda rep ->-  Lambda rep ->-  [SubExp] ->-  [VName] ->-  DistNestT rep m (Maybe (Stms rep))-regularSegmentedRedomapKernel nest perm cs segment_size comm lam map_lam nes arrs = do-  mk_lvl <- asks distSegLevel-  isSegmentedOp nest perm (freeIn lam) (freeIn map_lam) nes [] $-    \pat ispace inps nes' _ -> do-      let red_op = SegBinOp comm lam nes' mempty-      lvl <- mk_lvl (segment_size : map snd ispace) "segred" $ NoRecommendation SegNoVirt-      addStms-        =<< traverse renameStm-        =<< segRed lvl pat cs segment_size [red_op] map_lam arrs ispace inps--isSegmentedOp ::-  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>-  KernelNest ->-  [Int] ->-  Names ->-  Names ->-  [SubExp] ->-  [VName] ->-  ( Pat Type ->-    [(VName, SubExp)] ->-    [KernelInput] ->-    [SubExp] ->-    [VName] ->-    BuilderT rep m ()-  ) ->-  DistNestT rep m (Maybe (Stms rep))-isSegmentedOp nest perm free_in_op _free_in_fold_op nes arrs m = runMaybeT $ do-  -- We must verify that array inputs to the operation are inputs to-  -- the outermost loop nesting or free in the loop nest.  Nothing-  -- free in the op may be bound by the nest.  Furthermore, the-  -- neutral elements must be free in the loop nest.-  ---  -- We must summarise any names from free_in_op that are bound in the-  -- nest, and describe how to obtain them given segment indices.--  let bound_by_nest = boundInKernelNest nest--  (ispace, kernel_inps) <- flatKernel nest--  when (free_in_op `namesIntersect` bound_by_nest) $-    fail "Non-fold lambda uses nest-bound parameters."--  let indices = map fst ispace--      prepareNe (Var v)-        | v `nameIn` bound_by_nest =-            fail "Neutral element bound in nest"-      prepareNe ne = pure ne--      prepareArr arr =-        case find ((== arr) . kernelInputName) kernel_inps of-          Just inp-            | kernelInputIndices inp == map Var indices ->-                pure $ pure $ kernelInputArray inp-          Nothing-            | arr `notNameIn` bound_by_nest ->-                -- This input is something that is free inside-                -- the loop nesting. We will have to replicate-                -- it.-                pure $-                  letExp-                    (baseName arr <> "_repd")-                    (BasicOp $ Replicate (Shape $ map snd ispace) $ Var arr)-          _ ->-            fail "Input not free, perfectly mapped, or outermost."--  nes' <- mapM prepareNe nes--  mk_arrs <- mapM prepareArr arrs--  lift $-    liftInner $-      runBuilderT'_ $ do-        nested_arrs <- sequence mk_arrs--        let pat =-              Pat . rearrangeShape perm $-                patElems $-                  loopNestingPat $-                    fst nest--        m pat ispace kernel_inps nes' nested_arrs--permutationAndMissing :: Pat Type -> Result -> Maybe ([Int], [PatElem Type])-permutationAndMissing (Pat pes) res = do-  let (_used, unused) =-        partition ((`nameIn` freeIn res) . patElemName) pes-      res' = map resSubExp res-      res_expanded = res' ++ map (Var . patElemName) unused-  perm <- map (Var . patElemName) pes `isPermutationOf` res_expanded-  pure (perm, unused)---- Add extra pattern elements to every kernel nesting level.-expandKernelNest ::-  (MonadFreshNames m) => [PatElem Type] -> KernelNest -> m KernelNest-expandKernelNest pes (outer_nest, inner_nests) = do-  let outer_size =-        loopNestingWidth outer_nest-          : map loopNestingWidth inner_nests-      inner_sizes = tails $ map loopNestingWidth inner_nests-  outer_nest' <- expandWith outer_nest outer_size-  inner_nests' <- zipWithM expandWith inner_nests inner_sizes-  pure (outer_nest', inner_nests')-  where-    expandWith nest dims = do-      pes' <- mapM (expandPatElemWith dims) pes-      pure-        nest-          { loopNestingPat =-              Pat $ patElems (loopNestingPat nest) <> pes'-          }--    expandPatElemWith dims pe = do-      name <- newName $ patElemName pe-      pure-        pe-          { patElemName = name,-            patElemDec = patElemType pe `arrayOfShape` Shape dims-          }--kernelOrNot ::-  (MonadFreshNames m, DistRep rep) =>-  Certs ->-  Stm SOACS ->-  DistAcc rep ->-  PostStms rep ->-  DistAcc rep ->-  Maybe (Stms rep) ->-  DistNestT rep m (DistAcc rep)-kernelOrNot cs stm acc _ _ Nothing =-  addStmToAcc (certify cs stm) acc-kernelOrNot cs _ _ kernels acc' (Just stms) = do-  addPostStms kernels-  postStm $ fmap (certify cs) stms-  pure acc'--distributeMap ::-  (MonadFreshNames m, LocalScope rep m, DistRep rep) =>-  MapLoop ->-  DistAcc rep ->-  DistNestT rep m (DistAcc rep)-distributeMap (MapLoop pat aux w lam arrs) acc =-  distribute-    =<< mapNesting-      pat-      aux-      w-      lam-      arrs-      (distribute =<< distributeMapBodyStms acc' lam_stms)-  where-    acc' =-      DistAcc-        { distTargets =-            pushInnerTarget-              (pat, bodyResult $ lambdaBody lam)-              $ distTargets acc,-          distStms = mempty-        }--    lam_stms = bodyStms $ lambdaBody lam
− src/Futhark/Pass/ExtractKernels/Distribution.hs 
@@ -1,584 +0,0 @@-{-# LANGUAGE LambdaCase #-}-{-# LANGUAGE TypeFamilies #-}--module Futhark.Pass.ExtractKernels.Distribution-  ( Target,-    Targets,-    ppTargets,-    singleTarget,-    outerTarget,-    innerTarget,-    pushInnerTarget,-    popInnerTarget,-    targetsScope,-    LoopNesting (..),-    ppLoopNesting,-    scopeOfLoopNesting,-    Nesting (..),-    Nestings,-    ppNestings,-    letBindInInnerNesting,-    singleNesting,-    pushInnerNesting,-    KernelNest,-    ppKernelNest,-    newKernel,-    innermostKernelNesting,-    pushKernelNesting,-    pushInnerKernelNesting,-    scopeOfKernelNest,-    kernelNestLoops,-    kernelNestWidths,-    boundInKernelNest,-    boundInKernelNests,-    flatKernel,-    constructKernel,-    tryDistribute,-    tryDistributeStm,-  )-where--import Control.Monad-import Control.Monad.RWS.Strict-import Control.Monad.Trans.Maybe-import Data.Bifunctor (second)-import Data.Foldable-import Data.List (elemIndex, sortOn)-import Data.Map.Strict qualified as M-import Data.Maybe-import Futhark.IR-import Futhark.IR.SegOp-import Futhark.MonadFreshNames-import Futhark.Pass.ExtractKernels.BlockedKernel-  ( DistRep,-    KernelInput (..),-    MkSegLevel,-    mapKernel,-    readKernelInput,-  )-import Futhark.Tools-import Futhark.Transform.Rename-import Futhark.Util-import Futhark.Util.Log--type Target = (Pat Type, Result)---- | First pair element is the very innermost ("current") target.  In--- the list, the outermost target comes first.  Invariant: Every--- element of a pattern must be present as the result of the--- immediately enclosing target.  This is ensured by 'pushInnerTarget'--- by removing unused pattern elements.-data Targets = Targets-  { _innerTarget :: Target,-    _outerTargets :: [Target]-  }--ppTargets :: Targets -> String-ppTargets (Targets target targets) =-  unlines $ map ppTarget $ targets ++ [target]-  where-    ppTarget (pat, res) = prettyString pat ++ " <- " ++ prettyString res--singleTarget :: Target -> Targets-singleTarget = flip Targets []--outerTarget :: Targets -> Target-outerTarget (Targets inner_target []) = inner_target-outerTarget (Targets _ (outer_target : _)) = outer_target--innerTarget :: Targets -> Target-innerTarget (Targets inner_target _) = inner_target--pushOuterTarget :: Target -> Targets -> Targets-pushOuterTarget target (Targets inner_target targets) =-  Targets inner_target (target : targets)--pushInnerTarget :: Target -> Targets -> Targets-pushInnerTarget (pat, res) (Targets inner_target targets) =-  Targets (pat', res') (targets ++ [inner_target])-  where-    (pes', res') = unzip $ filter (used . fst) $ zip (patElems pat) res-    pat' = Pat pes'-    inner_used = freeIn $ snd inner_target-    used pe = patElemName pe `nameIn` inner_used--popInnerTarget :: Targets -> Maybe (Target, Targets)-popInnerTarget (Targets t ts) =-  case reverse ts of-    x : xs -> Just (t, Targets x $ reverse xs)-    [] -> Nothing--targetScope :: (DistRep rep) => Target -> Scope rep-targetScope = scopeOfPat . fst--targetsScope :: (DistRep rep) => Targets -> Scope rep-targetsScope (Targets t ts) = mconcat $ map targetScope $ t : ts--data LoopNesting = MapNesting-  { loopNestingPat :: Pat Type,-    loopNestingAux :: StmAux (),-    loopNestingWidth :: SubExp,-    loopNestingParamsAndArrs :: [(Param Type, VName)]-  }-  deriving (Show)--scopeOfLoopNesting :: (LParamInfo rep ~ Type) => LoopNesting -> Scope rep-scopeOfLoopNesting = scopeOfLParams . map fst . loopNestingParamsAndArrs--ppLoopNesting :: LoopNesting -> String-ppLoopNesting (MapNesting _ _ _ params_and_arrs) =-  prettyString (map fst params_and_arrs)-    ++ " <- "-    ++ prettyString (map snd params_and_arrs)--loopNestingParams :: LoopNesting -> [Param Type]-loopNestingParams = map fst . loopNestingParamsAndArrs--instance FreeIn LoopNesting where-  freeIn' (MapNesting pat aux w params_and_arrs) =-    freeIn' pat <> freeIn' aux <> freeIn' w <> freeIn' params_and_arrs--data Nesting = Nesting-  { nestingLetBound :: Names,-    nestingLoop :: LoopNesting-  }-  deriving (Show)--letBindInNesting :: Names -> Nesting -> Nesting-letBindInNesting newnames (Nesting oldnames loop) =-  Nesting (oldnames <> newnames) loop--- ^ First pair element is the very innermost ("current") nest.  In--- the list, the outermost nest comes first.--type Nestings = (Nesting, [Nesting])--ppNestings :: Nestings -> String-ppNestings (nesting, nestings) =-  unlines $ map ppNesting $ nestings ++ [nesting]-  where-    ppNesting (Nesting _ loop) = ppLoopNesting loop--singleNesting :: Nesting -> Nestings-singleNesting = (,[])--pushInnerNesting :: Nesting -> Nestings -> Nestings-pushInnerNesting nesting (inner_nesting, nestings) =-  (nesting, nestings ++ [inner_nesting])---- | Both parameters and let-bound.-boundInNesting :: Nesting -> Names-boundInNesting nesting =-  namesFromList (map paramName (loopNestingParams loop))-    <> nestingLetBound nesting-  where-    loop = nestingLoop nesting--letBindInInnerNesting :: Names -> Nestings -> Nestings-letBindInInnerNesting names (nest, nestings) =-  (letBindInNesting names nest, nestings)---- | Note: first element is *outermost* nesting.  This is different--- from the similar types elsewhere!-type KernelNest = (LoopNesting, [LoopNesting])--ppKernelNest :: KernelNest -> String-ppKernelNest (nesting, nestings) =-  unlines $ map ppLoopNesting $ nesting : nestings---- | Retrieve the innermost kernel nesting.-innermostKernelNesting :: KernelNest -> LoopNesting-innermostKernelNesting (nest, nests) =-  fromMaybe nest $ maybeHead $ reverse nests---- | Add new outermost nesting, pushing the current outermost to the--- list, also taking care to swap patterns if necessary.-pushKernelNesting :: Target -> LoopNesting -> KernelNest -> KernelNest-pushKernelNesting target newnest (nest, nests) =-  ( fixNestingPatOrder newnest target (loopNestingPat nest),-    nest : nests-  )---- | Add new innermost nesting, pushing the current outermost to the--- list.  It is important that the 'Target' has the right order--- (non-permuted compared to what is expected by the outer nests).-pushInnerKernelNesting :: Target -> LoopNesting -> KernelNest -> KernelNest-pushInnerKernelNesting target newnest (nest, nests) =-  (nest, nests ++ [fixNestingPatOrder newnest target (loopNestingPat innermost)])-  where-    innermost = case reverse nests of-      [] -> nest-      n : _ -> n--fixNestingPatOrder :: LoopNesting -> Target -> Pat Type -> LoopNesting-fixNestingPatOrder nest (_, res) inner_pat =-  nest {loopNestingPat = basicPat pat'}-  where-    pat = loopNestingPat nest-    pat' = map fst fixed_target-    fixed_target = sortOn posInInnerPat $ zip (patIdents pat) res-    posInInnerPat (_, SubExpRes _ (Var v)) = fromMaybe 0 $ elemIndex v $ patNames inner_pat-    posInInnerPat _ = 0--newKernel :: LoopNesting -> KernelNest-newKernel nest = (nest, [])--kernelNestLoops :: KernelNest -> [LoopNesting]-kernelNestLoops (loop, loops) = loop : loops--scopeOfKernelNest :: (LParamInfo rep ~ Type) => KernelNest -> Scope rep-scopeOfKernelNest = foldMap scopeOfLoopNesting . kernelNestLoops--boundInKernelNest :: KernelNest -> Names-boundInKernelNest = mconcat . boundInKernelNests--boundInKernelNests :: KernelNest -> [Names]-boundInKernelNests =-  map (namesFromList . map (paramName . fst) . loopNestingParamsAndArrs)-    . kernelNestLoops--kernelNestWidths :: KernelNest -> [SubExp]-kernelNestWidths = map loopNestingWidth . kernelNestLoops--constructKernel ::-  (DistRep rep, MonadFreshNames m, LocalScope rep m) =>-  MkSegLevel rep m ->-  KernelNest ->-  Body rep ->-  m (Stm rep, Stms rep)-constructKernel mk_lvl kernel_nest inner_body = runBuilderT' $ do-  (ispace, inps) <- flatKernel kernel_nest-  let aux = loopNestingAux first_nest-      ispace_scope = M.fromList $ map ((,IndexName Int64) . fst) ispace-      pat = loopNestingPat first_nest-      rts = map (stripArray (length ispace)) $ patTypes pat--  inner_body' <- fmap (uncurry (flip (Body ()))) $-    runBuilder . localScope ispace_scope $ do-      mapM_ readKernelInput $ filter inputIsUsed inps-      res <- bodyBind inner_body-      forM res $ \(SubExpRes cs se) -> pure $ Returns ResultMaySimplify cs se--  (segop, aux_stms) <- lift $ mapKernel mk_lvl ispace [] rts inner_body'--  addStms aux_stms--  pure $ Let pat aux $ Op $ segOp segop-  where-    first_nest = fst kernel_nest-    inputIsUsed input = kernelInputName input `nameIn` freeIn inner_body---- | Flatten a kernel nesting to:------  (1) The index space.------  (2) The kernel inputs - note that some of these may be unused.-flatKernel ::-  (MonadFreshNames m) =>-  KernelNest ->-  m ([(VName, SubExp)], [KernelInput])-flatKernel (MapNesting _ _ nesting_w params_and_arrs, []) = do-  i <- newVName "gtid"-  let inps =-        [ KernelInput pname ptype arr [Var i]-        | (Param _ pname ptype, arr) <- params_and_arrs-        ]-  pure ([(i, nesting_w)], inps)-flatKernel (MapNesting _ _ nesting_w params_and_arrs, nest : nests) = do-  i <- newVName "gtid"-  (ispace, inps) <- flatKernel (nest, nests)--  let inps' = map fixupInput inps-      isParam inp =-        snd <$> find ((== kernelInputArray inp) . paramName . fst) params_and_arrs-      fixupInput inp-        | Just arr <- isParam inp =-            inp-              { kernelInputArray = arr,-                kernelInputIndices = Var i : kernelInputIndices inp-              }-        | otherwise =-            inp--  pure ((i, nesting_w) : ispace, extra_inps i <> inps')-  where-    extra_inps i =-      [ KernelInput pname ptype arr [Var i]-      | (Param _ pname ptype, arr) <- params_and_arrs-      ]---- | Description of distribution to do.-data DistributionBody = DistributionBody-  { distributionTarget :: Targets,-    distributionFreeInBody :: Names,-    distributionIdentityMap :: M.Map VName Ident,-    -- | Also related to avoiding identity mapping.-    distributionExpandTarget :: Target -> Target-  }--distributionInnerPat :: DistributionBody -> Pat Type-distributionInnerPat = fst . innerTarget . distributionTarget--distributionBodyFromStms ::-  (ASTRep rep) =>-  Targets ->-  Stms rep ->-  (DistributionBody, Result)-distributionBodyFromStms (Targets (inner_pat, inner_res) targets) stms =-  let bound_by_stms = namesFromList $ M.keys $ scopeOf stms-      (inner_pat', inner_res', inner_identity_map, inner_expand_target) =-        removeIdentityMappingGeneral bound_by_stms inner_pat inner_res-      free =-        (foldMap freeIn stms <> freeIn (map resCerts inner_res))-          `namesSubtract` bound_by_stms-   in ( DistributionBody-          { distributionTarget = Targets (inner_pat', inner_res') targets,-            distributionFreeInBody = free,-            distributionIdentityMap = inner_identity_map,-            distributionExpandTarget = inner_expand_target-          },-        inner_res'-      )--distributionBodyFromStm ::-  (ASTRep rep) =>-  Targets ->-  Stm rep ->-  (DistributionBody, Result)-distributionBodyFromStm targets stm =-  distributionBodyFromStms targets $ oneStm stm--createKernelNest ::-  forall rep m.-  (MonadFreshNames m, HasScope rep m) =>-  Nestings ->-  DistributionBody ->-  m (Maybe (Targets, KernelNest))-createKernelNest (inner_nest, nests) distrib_body = do-  let Targets target targets = distributionTarget distrib_body-  unless (length nests == length targets) $-    error $-      "Nests and targets do not match!\n"-        ++ "nests: "-        ++ ppNestings (inner_nest, nests)-        ++ "\ntargets:"-        ++ ppTargets (Targets target targets)-  runMaybeT $ fmap prepare $ recurse $ zip nests targets-  where-    prepare (x, _, z) = (z, x)-    bound_in_nest = mconcat $ map boundInNesting $ inner_nest : nests-    distributableType =-      (== mempty) . namesIntersection bound_in_nest . freeIn . arrayDims--    distributeAtNesting ::-      Nesting ->-      Pat Type ->-      (LoopNesting -> KernelNest, Names) ->-      M.Map VName Ident ->-      [Ident] ->-      (Target -> Targets) ->-      MaybeT m (KernelNest, Names, Targets)-    distributeAtNesting-      (Nesting nest_let_bound nest)-      pat-      (add_to_kernel, free_in_kernel)-      identity_map-      inner_returned_arrs-      addTarget = do-        let nest'@(MapNesting _ aux w params_and_arrs) =-              removeUnusedNestingParts free_in_kernel nest-            (params, arrs) = unzip params_and_arrs-            param_names = namesFromList $ map paramName params-            free_in_kernel' =-              (freeIn nest' <> free_in_kernel) `namesSubtract` param_names-            required_from_nest =-              free_in_kernel' `namesIntersection` nest_let_bound--        required_from_nest_idents <--          forM (namesToList required_from_nest) $ \name -> do-            t <- lift $ lookupType name-            pure $ Ident name t--        (free_params, free_arrs, bind_in_target) <--          fmap unzip3 $-            forM (inner_returned_arrs ++ required_from_nest_idents) $-              \(Ident pname ptype) ->-                case M.lookup pname identity_map of-                  Nothing -> do-                    arr <--                      newIdent (baseName pname <> "_r") $ arrayOfRow ptype w-                    pure-                      ( Param mempty pname ptype,-                        arr,-                        True-                      )-                  Just arr ->-                    pure-                      ( Param mempty pname ptype,-                        arr,-                        False-                      )--        let free_arrs_pat =-              basicPat $ map snd $ filter fst $ zip bind_in_target free_arrs-            free_params_pat =-              map snd $ filter fst $ zip bind_in_target free_params--            (actual_params, actual_arrs) =-              ( params ++ free_params,-                arrs ++ map identName free_arrs-              )-            actual_param_names =-              namesFromList $ map paramName actual_params--            nest'' =-              removeUnusedNestingParts free_in_kernel $-                MapNesting pat aux w $-                  zip actual_params actual_arrs--            free_in_kernel'' =-              (freeIn nest'' <> free_in_kernel) `namesSubtract` actual_param_names--        unless-          ( all (distributableType . paramType) $-              loopNestingParams nest''-          )-          $ fail "Would induce irregular array"-        pure-          ( add_to_kernel nest'',-            free_in_kernel'',-            addTarget (free_arrs_pat, varsRes $ map paramName free_params_pat)-          )--    recurse :: [(Nesting, Target)] -> MaybeT m (KernelNest, Names, Targets)-    recurse [] =-      distributeAtNesting-        inner_nest-        (distributionInnerPat distrib_body)-        ( newKernel,-          distributionFreeInBody distrib_body `namesIntersection` bound_in_nest-        )-        (distributionIdentityMap distrib_body)-        []-        $ singleTarget . distributionExpandTarget distrib_body-    recurse ((nest, (pat, res)) : nests') = do-      (kernel@(outer, _), kernel_free, kernel_targets) <- recurse nests'--      let (pat', res', identity_map, expand_target) =-            removeIdentityMappingFromNesting-              (namesFromList $ patNames $ loopNestingPat outer)-              pat-              res--      distributeAtNesting-        nest-        pat'-        ( \k -> pushKernelNesting (pat', res') k kernel,-          kernel_free-        )-        identity_map-        (patIdents $ fst $ outerTarget kernel_targets)-        ((`pushOuterTarget` kernel_targets) . expand_target)--removeUnusedNestingParts :: Names -> LoopNesting -> LoopNesting-removeUnusedNestingParts used (MapNesting pat aux w params_and_arrs) =-  MapNesting pat aux w $ zip used_params used_arrs-  where-    (params, arrs) = unzip params_and_arrs-    (used_params, used_arrs) =-      unzip $ filter ((`nameIn` used) . paramName . fst) $ zip params arrs--removeIdentityMappingGeneral ::-  Names ->-  Pat Type ->-  Result ->-  ( Pat Type,-    Result,-    M.Map VName Ident,-    Target -> Target-  )-removeIdentityMappingGeneral bound pat res =-  let (identities, not_identities) =-        mapEither isIdentity $ zip (patElems pat) res-      (not_identity_patElems, not_identity_res) = unzip not_identities-      (identity_patElems, identity_res) = unzip identities-      expandTarget (tpat, tres) =-        ( Pat $ patElems tpat ++ identity_patElems,-          tres ++ map (uncurry SubExpRes . second Var) identity_res-        )-      identity_map =-        M.fromList $ zip (map snd identity_res) $ map patElemIdent identity_patElems-   in ( Pat not_identity_patElems,-        not_identity_res,-        identity_map,-        expandTarget-      )-  where-    isIdentity (patElem, SubExpRes _ (Var v))-      | v `notNameIn` bound = Left (patElem, (mempty, v))-    isIdentity x = Right x--removeIdentityMappingFromNesting ::-  Names ->-  Pat Type ->-  Result ->-  ( Pat Type,-    Result,-    M.Map VName Ident,-    Target -> Target-  )-removeIdentityMappingFromNesting bound_in_nesting pat res =-  let (pat', res', identity_map, expand_target) =-        removeIdentityMappingGeneral bound_in_nesting pat res-   in (pat', res', identity_map, expand_target)--tryDistribute ::-  ( DistRep rep,-    MonadFreshNames m,-    LocalScope rep m,-    MonadLogger m-  ) =>-  MkSegLevel rep m ->-  Nestings ->-  Targets ->-  Stms rep ->-  m (Maybe (Targets, Stms rep))-tryDistribute _ _ targets stms-  | null stms =-      -- No point in distributing an empty kernel.-      pure $ Just (targets, mempty)-tryDistribute mk_lvl nest targets stms =-  createKernelNest nest dist_body-    >>= \case-      Just (targets', distributed) -> do-        (kernel_stm, w_stms) <--          localScope (targetsScope targets') $-            constructKernel mk_lvl distributed $-              mkBody stms inner_body_res-        distributed' <- renameStm kernel_stm-        logMsg $-          "distributing\n"-            ++ unlines (map prettyString $ stmsToList stms)-            ++ prettyString (snd $ innerTarget targets)-            ++ "\nas\n"-            ++ prettyString distributed'-            ++ "\ndue to targets\n"-            ++ ppTargets targets-            ++ "\nand with new targets\n"-            ++ ppTargets targets'-        pure $ Just (targets', w_stms <> oneStm distributed')-      Nothing ->-        pure Nothing-  where-    (dist_body, inner_body_res) = distributionBodyFromStms targets stms--tryDistributeStm ::-  (MonadFreshNames m, HasScope t m, ASTRep rep) =>-  Nestings ->-  Targets ->-  Stm rep ->-  m (Maybe (Result, Targets, KernelNest))-tryDistributeStm nest targets stm =-  fmap addRes <$> createKernelNest nest dist_body-  where-    (dist_body, res) = distributionBodyFromStm targets stm-    addRes (targets', kernel_nest) = (res, targets', kernel_nest)
− src/Futhark/Pass/ExtractKernels/ISRWIM.hs 
@@ -1,195 +0,0 @@-{-# LANGUAGE TypeFamilies #-}---- | Interchanging scans with inner maps.-module Futhark.Pass.ExtractKernels.ISRWIM-  ( iswim,-    irwim,-    rwimPossible,-  )-where--import Control.Arrow (first)-import Control.Monad-import Futhark.IR.SOACS-import Futhark.MonadFreshNames-import Futhark.Tools---- | Interchange Scan With Inner Map. Tries to turn a @scan(map)@ into a--- @map(scan)-iswim ::-  (MonadBuilder m, Rep m ~ SOACS) =>-  Pat Type ->-  SubExp ->-  Lambda SOACS ->-  [(SubExp, VName)] ->-  Maybe (m ())-iswim res_pat w scan_fun scan_input-  | Just (map_pat, map_aux, map_w, map_fun) <- rwimPossible scan_fun = Just $ do-      let (accs, arrs) = unzip scan_input-      arrs' <- transposedArrays arrs-      accs' <- mapM (letExp "acc" . BasicOp . SubExp) accs--      let map_arrs' = accs' ++ arrs'-          (scan_acc_params, scan_elem_params) =-            splitAt (length arrs) $ lambdaParams scan_fun-          map_params =-            map removeParamOuterDim scan_acc_params-              ++ map (setParamOuterDimTo w) scan_elem_params-          map_rettype = map (setOuterDimTo w) $ lambdaReturnType scan_fun--          scan_params = lambdaParams map_fun-          scan_body = lambdaBody map_fun-          scan_rettype = lambdaReturnType map_fun-          scan_fun' = Lambda scan_params scan_rettype scan_body-          scan_input' =-            map (first Var) $-              uncurry zip $-                splitAt (length arrs') $-                  map paramName map_params-          (nes', scan_arrs) = unzip scan_input'--      scan_soac <- scanSOAC [Scan scan_fun' nes']-      let map_body =-            mkBody-              ( oneStm $-                  Let (setPatOuterDimTo w map_pat) (defAux ()) $-                    Op $-                      Screma w scan_arrs scan_soac-              )-              $ varsRes-              $ patNames map_pat-          map_fun' = Lambda map_params map_rettype map_body--      res_pat' <--        fmap basicPat $-          mapM (newIdent' (<> "_transposed") . transposeIdentType) $-            patIdents res_pat--      addStm . Let res_pat' map_aux . Op . Screma map_w map_arrs'-        =<< mapSOAC map_fun'--      forM_ (zip (patIdents res_pat) (patIdents res_pat')) $ \(to, from) -> do-        let perm = [1, 0] ++ [2 .. arrayRank (identType from) - 1]-        addStm $-          Let (basicPat [to]) (defAux ()) . BasicOp $-            Rearrange (identName from) perm-  | otherwise = Nothing---- | Interchange Reduce With Inner Map. Tries to turn a @reduce(map)@ into a--- @map(reduce)-irwim ::-  (MonadBuilder m, Rep m ~ SOACS) =>-  Pat Type ->-  SubExp ->-  Commutativity ->-  Lambda SOACS ->-  [(SubExp, VName)] ->-  Maybe (m ())-irwim res_pat w comm red_fun red_input-  | Just (map_pat, map_aux, map_w, map_fun) <- rwimPossible red_fun = Just $ do-      let (accs, arrs) = unzip red_input-      arrs' <- transposedArrays arrs-      -- FIXME?  Can we reasonably assume that the accumulator is a-      -- replicate?  We also assume that it is non-empty.-      let indexAcc (Var v) = do-            v_t <- lookupType v-            letSubExp "acc" $-              BasicOp $-                Index v $-                  fullSlice v_t [DimFix $ intConst Int64 0]-          indexAcc Constant {} =-            error "irwim: array accumulator is a constant."-      accs' <- mapM indexAcc accs--      let (_red_acc_params, red_elem_params) =-            splitAt (length arrs) $ lambdaParams red_fun-          map_rettype = map rowType $ lambdaReturnType red_fun-          map_params = map (setParamOuterDimTo w) red_elem_params--          red_params = lambdaParams map_fun-          red_body = lambdaBody map_fun-          red_rettype = lambdaReturnType map_fun-          red_fun' = Lambda red_params red_rettype red_body-          red_input' = zip accs' $ map paramName map_params-          red_pat = stripPatOuterDim map_pat--      map_body <--        case irwim red_pat w comm red_fun' red_input' of-          Nothing -> do-            reduce_soac <- reduceSOAC [Reduce comm red_fun' $ map fst red_input']-            pure-              $ mkBody-                ( oneStm $-                    Let red_pat (defAux ()) $-                      Op $-                        Screma w (map snd red_input') reduce_soac-                )-              $ varsRes-              $ patNames map_pat-          Just m -> localScope (scopeOfLParams map_params) $ do-            map_body_stms <- collectStms_ m-            pure $ mkBody map_body_stms $ varsRes $ patNames map_pat--      let map_fun' = Lambda map_params map_rettype map_body--      addStm . Let res_pat map_aux . Op . Screma map_w arrs'-        =<< mapSOAC map_fun'-  | otherwise = Nothing---- | Does this reduce operator contain an inner map, and if so, what--- does that map look like?-rwimPossible ::-  Lambda SOACS ->-  Maybe (Pat Type, StmAux (), SubExp, Lambda SOACS)-rwimPossible fun-  | Body _ stms res <- lambdaBody fun,-    [stm] <- stmsToList stms, -- Body has a single binding-    map_pat <- stmPat stm,-    map Var (patNames map_pat) == map resSubExp res, -- Returned verbatim-    Op (Screma map_w map_arrs form) <- stmExp stm,-    Just map_fun <- isMapSOAC form,-    map paramName (lambdaParams fun) == map_arrs =-      Just (map_pat, stmAux stm, map_w, map_fun)-  | otherwise =-      Nothing--transposedArrays :: (MonadBuilder m) => [VName] -> m [VName]-transposedArrays arrs = forM arrs $ \arr -> do-  t <- lookupType arr-  let perm = [1, 0] ++ [2 .. arrayRank t - 1]-  letExp (baseName arr) $ BasicOp $ Rearrange arr perm--removeParamOuterDim :: LParam SOACS -> LParam SOACS-removeParamOuterDim param =-  let t = rowType $ paramType param-   in param {paramDec = t}--setParamOuterDimTo :: SubExp -> LParam SOACS -> LParam SOACS-setParamOuterDimTo w param =-  let t = setOuterDimTo w $ paramType param-   in param {paramDec = t}--setIdentOuterDimTo :: SubExp -> Ident -> Ident-setIdentOuterDimTo w ident =-  let t = setOuterDimTo w $ identType ident-   in ident {identType = t}--setOuterDimTo :: SubExp -> Type -> Type-setOuterDimTo w t =-  arrayOfRow (rowType t) w--setPatOuterDimTo :: SubExp -> Pat Type -> Pat Type-setPatOuterDimTo w pat =-  basicPat $ map (setIdentOuterDimTo w) $ patIdents pat--transposeIdentType :: Ident -> Ident-transposeIdentType ident =-  ident {identType = transposeType $ identType ident}--stripIdentOuterDim :: Ident -> Ident-stripIdentOuterDim ident =-  ident {identType = rowType $ identType ident}--stripPatOuterDim :: Pat Type -> Pat Type-stripPatOuterDim pat =-  basicPat $ map stripIdentOuterDim $ patIdents pat
− src/Futhark/Pass/ExtractKernels/Interchange.hs 
@@ -1,357 +0,0 @@-{-# LANGUAGE TypeFamilies #-}---- | It is well known that fully parallel loops can always be--- interchanged inwards with a sequential loop.  This module--- implements that transformation.------ This is also where we implement loop-switching (for branches),--- which is semantically similar to interchange.-module Futhark.Pass.ExtractKernels.Interchange-  ( SeqLoop (..),-    interchangeLoops,-    Branch (..),-    interchangeBranch,-    WithAccStm (..),-    interchangeWithAcc,-  )-where--import Control.Monad-import Data.List (find)-import Data.Maybe-import Futhark.IR.SOACS-import Futhark.MonadFreshNames-import Futhark.Pass.ExtractKernels.Distribution-  ( KernelNest,-    LoopNesting (..),-    kernelNestLoops,-    scopeOfKernelNest,-  )-import Futhark.Tools-import Futhark.Transform.Rename-import Futhark.Util (splitFromEnd)---- | An encoding of a sequential do-loop with no existential context,--- alongside its result pattern.-data SeqLoop-  = SeqLoop [Int] (Pat Type) [(FParam SOACS, SubExp)] LoopForm (Body SOACS)--loopPerm :: SeqLoop -> [Int]-loopPerm (SeqLoop perm _ _ _ _) = perm--seqLoopStm :: SeqLoop -> Stm SOACS-seqLoopStm (SeqLoop _ pat merge form body) =-  Let pat (defAux ()) $ Loop merge form body--interchangeLoop ::-  (MonadBuilder m, Rep m ~ SOACS) =>-  (VName -> Maybe VName) ->-  SeqLoop ->-  LoopNesting ->-  m SeqLoop-interchangeLoop-  isMapParameter-  (SeqLoop perm loop_pat merge form body)-  (MapNesting pat aux w params_and_arrs) = do-    merge_expanded <--      localScope (scopeOfLParams $ map fst params_and_arrs) $-        mapM expand merge--    let loop_pat_expanded =-          Pat $ map expandPatElem $ patElems loop_pat-        new_params =-          [Param attrs pname $ fromDecl ptype | (Param attrs pname ptype, _) <- merge]-        new_arrs = map (paramName . fst) merge_expanded-        rettype = map rowType $ patTypes loop_pat_expanded--    -- If the map consumes something that is bound outside the loop-    -- (i.e. is not a merge parameter), we have to copy() it.  As a-    -- small simplification, we just remove the parameter outright if-    -- it is not used anymore.  This might happen if the parameter was-    -- used just as the inital value of a merge parameter.-    ((params', arrs'), pre_copy_stms) <--      runBuilder $-        localScope (scopeOfLParams new_params) $-          unzip . catMaybes <$> mapM copyOrRemoveParam params_and_arrs--    let lam = Lambda (params' <> new_params) rettype body-    map_stm <--      Let loop_pat_expanded aux . Op . Screma w (arrs' <> new_arrs)-        <$> mapSOAC lam-    let res = varsRes $ patNames loop_pat_expanded-        pat' = Pat $ rearrangeShape perm $ patElems pat--    pure $-      SeqLoop perm pat' merge_expanded form $-        mkBody (pre_copy_stms <> oneStm map_stm) res-    where-      free_in_body = freeIn body--      copyOrRemoveParam (param, arr)-        | paramName param `notNameIn` free_in_body =-            pure Nothing-        | otherwise =-            pure $ Just (param, arr)--      expandedInit _ (Var v)-        | Just arr <- isMapParameter v =-            pure $ Var arr-      expandedInit param_name se =-        letSubExp (param_name <> "_expanded_init") $-          BasicOp $-            Replicate (Shape [w]) se--      expand (merge_param, merge_init) = do-        expanded_param <--          newParam (param_name <> "_expanded") $-            -- FIXME: Unique here is a hack to make sure the copy from-            -- makeCopyInitial is not prematurely simplified away.-            -- It'd be better to fix this somewhere else...-            arrayOf (paramDeclType merge_param) (Shape [w]) Unique-        expanded_init <- expandedInit param_name merge_init-        pure (expanded_param, expanded_init)-        where-          param_name = baseName $ paramName merge_param--      expandPatElem (PatElem name t) =-        PatElem name $ arrayOfRow t w---- We need to copy some initial arguments because otherwise the result--- of the loop might alias the input (if the number of iterations is--- 0), which is a problem if the result is consumed.-maybeCopyInitial ::-  (MonadBuilder m) =>-  (VName -> Bool) ->-  SeqLoop ->-  m SeqLoop-maybeCopyInitial isMapInput (SeqLoop perm loop_pat merge form body) =-  SeqLoop perm loop_pat <$> mapM f merge <*> pure form <*> pure body-  where-    f (p, Var arg)-      | isMapInput arg,-        Array {} <- paramType p =-          (p,)-            <$> letSubExp-              (baseName (paramName p) <> "_inter_copy")-              (BasicOp $ Replicate mempty $ Var arg)-    f (p, arg) =-      pure (p, arg)--manifestMaps ::-  (MonadFreshNames m) =>-  [LoopNesting] ->-  [VName] ->-  Stms SOACS ->-  m ([VName], Stms SOACS)-manifestMaps [] res stms = pure (res, stms)-manifestMaps (n : ns) res stms = do-  (res', stms') <- manifestMaps ns res stms-  let (params, arrs) = unzip $ loopNestingParamsAndArrs n-      lam =-        Lambda-          params-          (map rowType $ patTypes (loopNestingPat n))-          (mkBody stms' $ varsRes res')-  st <--    oneStm-      . Let (loopNestingPat n) (loopNestingAux n)-      . Op-      . Screma (loopNestingWidth n) arrs-      <$> mapSOAC lam-  pure-    (patNames $ loopNestingPat n, st)---- | Given a (parallel) map nesting and an inner sequential loop, move--- the maps inside the sequential loop.  The result is several--- statements - one of these will be the loop, which will then contain--- statements with @map@ expressions.-interchangeLoops ::-  (MonadFreshNames m, HasScope SOACS m) =>-  KernelNest ->-  SeqLoop ->-  m (Stms SOACS)-interchangeLoops full_nest = recurse (kernelNestLoops full_nest)-  where-    recurse nest loop-      | (ns, [n]) <- splitFromEnd 1 nest = do-          let isMapParameter v =-                snd <$> find ((== v) . paramName . fst) (loopNestingParamsAndArrs n)-              isMapInput v =-                v `elem` map snd (loopNestingParamsAndArrs n)-          (loop', stms) <--            runBuilder . localScope (scopeOfKernelNest full_nest) $-              maybeCopyInitial isMapInput-                =<< interchangeLoop isMapParameter loop n--          -- Only safe to continue interchanging if we didn't need to add-          -- any new statements; otherwise we manifest the remaining nests-          -- as Maps and hand them back to the flattener.-          if null stms-            then recurse ns loop'-            else-              let loop_stm = seqLoopStm loop'-                  names = rearrangeShape (loopPerm loop') (patNames (stmPat loop_stm))-               in snd <$> manifestMaps ns names (stms <> oneStm loop_stm)-      | otherwise = pure $ oneStm $ seqLoopStm loop---- | An encoding of a branch with alongside its result pattern.-data Branch-  = Branch [Int] (Pat Type) [SubExp] [Case (Body SOACS)] (Body SOACS) (MatchDec (BranchType SOACS))--branchStm :: Branch -> Stm SOACS-branchStm (Branch _ pat cond cases defbody ret) =-  Let pat (defAux ()) $ Match cond cases defbody ret--interchangeBranch1 ::-  (MonadFreshNames m, HasScope SOACS m) =>-  Branch ->-  LoopNesting ->-  m Branch-interchangeBranch1-  (Branch perm branch_pat cond cases defbody (MatchDec ret if_sort))-  (MapNesting pat aux w params_and_arrs) = do-    let ret' = map (`arrayOfRow` Free w) ret-        pat' = Pat $ rearrangeShape perm $ patElems pat--        (params, arrs) = unzip params_and_arrs-        lam_ret = rearrangeShape perm $ map rowType $ patTypes pat--        branch_pat' =-          Pat $ map (fmap (`arrayOfRow` w)) $ patElems branch_pat--        mkBranch branch = (renameBody =<<) $ runBodyBuilder $ do-          let lam = Lambda params lam_ret branch-          addStm . Let branch_pat' aux . Op . Screma w arrs =<< mapSOAC lam-          pure $ varsRes $ patNames branch_pat'--    cases' <- mapM (traverse mkBranch) cases-    defbody' <- mkBranch defbody-    pure . Branch [0 .. patSize pat - 1] pat' cond cases' defbody' $-      MatchDec ret' if_sort---- | Given a (parallel) map nesting and an inner branch, move the maps--- inside the branch.  The result is the resulting branch expression,--- which will then contain statements with @map@ expressions.-interchangeBranch ::-  (MonadFreshNames m, HasScope SOACS m) =>-  KernelNest ->-  Branch ->-  m (Stm SOACS)-interchangeBranch nest loop =-  branchStm <$> foldM interchangeBranch1 loop (reverse $ kernelNestLoops nest)---- | An encoding of a WithAcc with alongside its result pattern.-data WithAccStm-  = WithAccStm [Int] (Pat Type) [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))] (Lambda SOACS)--withAccStm :: WithAccStm -> Stm SOACS-withAccStm (WithAccStm _ pat inputs lam) =-  Let pat (defAux ()) $ WithAcc inputs lam--interchangeWithAcc1 ::-  (MonadFreshNames m, LocalScope SOACS m) =>-  WithAccStm ->-  LoopNesting ->-  m WithAccStm-interchangeWithAcc1-  (WithAccStm perm _withacc_pat inputs acc_lam)-  (MapNesting map_pat map_aux w params_and_arrs) = do-    inputs' <- mapM onInput inputs-    lam_params' <- newAccLamParams $ lambdaParams acc_lam-    iota_p <- newParam "iota_p" $ Prim int64-    acc_lam' <- trLam (Var (paramName iota_p)) <=< runLambdaBuilder lam_params' $ do-      let acc_params = drop (length inputs) lam_params'-          orig_acc_params = drop (length inputs) $ lambdaParams acc_lam-      iota_w <--        letExp "acc_inter_iota" . BasicOp $-          Iota w (intConst Int64 0) (intConst Int64 1) Int64-      let (params, arrs) = unzip params_and_arrs-          maplam_ret = lambdaReturnType acc_lam-          maplam = Lambda (iota_p : orig_acc_params ++ params) maplam_ret (lambdaBody acc_lam)-      auxing map_aux-        . fmap subExpsRes-        . letTupExp' "withacc_inter"-        . Op-        . Screma w (iota_w : map paramName acc_params ++ arrs)-        =<< mapSOAC maplam-    let pat = Pat $ rearrangeShape perm $ patElems map_pat-    pure $ WithAccStm perm pat inputs' acc_lam'-    where-      newAccLamParams ps = do-        let (cert_ps, acc_ps) = splitAt (length ps `div` 2) ps-        -- Should not rename the certificates.-        acc_ps' <- forM acc_ps $ \(Param attrs v t) ->-          Param attrs <$> newName v <*> pure t-        pure $ cert_ps <> acc_ps'--      num_accs = length inputs-      acc_certs = map paramName $ take num_accs $ lambdaParams acc_lam-      onArr v =-        pure . maybe v snd $-          find ((== v) . paramName . fst) params_and_arrs-      onInput (shape, arrs, op) =-        (Shape [w] <> shape,,) <$> mapM onArr arrs <*> traverse onOp op--      onOp (op_lam, nes) = do-        -- We need to add an additional index parameter because we are-        -- extending the index space of the accumulator.-        idx_p <- newParam "idx" $ Prim int64-        pure (op_lam {lambdaParams = idx_p : lambdaParams op_lam}, nes)--      trType :: TypeBase shape u -> TypeBase shape u-      trType (Acc acc ispace ts u)-        | acc `elem` acc_certs =-            Acc acc (Shape [w] <> ispace) ts u-      trType t = t--      trParam :: Param (TypeBase shape u) -> Param (TypeBase shape u)-      trParam = fmap trType--      trLam i (Lambda params ret body) =-        localScope (scopeOfLParams params) $-          Lambda (map trParam params) (map trType ret) <$> trBody i body--      trBody i (Body dec stms res) =-        inScopeOf stms $ Body dec <$> traverse (trStm i) stms <*> pure res--      trStm i (Let pat aux e) =-        Let (fmap trType pat) aux <$> trExp i e--      trSOAC i = mapSOACM mapper-        where-          mapper =-            identitySOACMapper {mapOnSOACLambda = trLam i}--      trExp i (WithAcc acc_inputs lam) =-        WithAcc acc_inputs <$> trLam i lam-      trExp i (BasicOp (UpdateAcc safety acc is ses)) = do-        acc_t <- lookupType acc-        pure $ case acc_t of-          Acc cert _ _ _-            | cert `elem` acc_certs ->-                BasicOp $ UpdateAcc safety acc (i : is) ses-          _ ->-            BasicOp $ UpdateAcc safety acc is ses-      trExp i e = mapExpM mapper e-        where-          mapper =-            identityMapper-              { mapOnBody = \scope -> localScope scope . trBody i,-                mapOnRetType = pure . trType,-                mapOnBranchType = pure . trType,-                mapOnFParam = pure . trParam,-                mapOnLParam = pure . trParam,-                mapOnOp = trSOAC i-              }---- | Given a (parallel) map nesting and an inner withacc, move the--- maps inside the branch.  The result is the resulting withacc--- expression, which will then contain statements with @map@--- expressions.-interchangeWithAcc ::-  (MonadFreshNames m, LocalScope SOACS m) =>-  KernelNest ->-  WithAccStm ->-  m (Stm SOACS)-interchangeWithAcc nest withacc =-  withAccStm <$> foldM interchangeWithAcc1 withacc (reverse $ kernelNestLoops nest)
− src/Futhark/Pass/ExtractKernels/Intrablock.hs 
@@ -1,318 +0,0 @@-{-# LANGUAGE TypeFamilies #-}---- | Extract limited nested parallelism for execution inside--- individual kernel threadblocks.-module Futhark.Pass.ExtractKernels.Intrablock (intrablockParallelise) where--import Control.Monad-import Control.Monad.RWS-import Control.Monad.Trans.Maybe-import Data.Map.Strict qualified as M-import Data.Set qualified as S-import Futhark.Analysis.PrimExp.Convert-import Futhark.IR.GPU hiding (HistOp)-import Futhark.IR.GPU.Op qualified as GPU-import Futhark.IR.SOACS-import Futhark.MonadFreshNames-import Futhark.Pass.ExtractKernels.BlockedKernel-import Futhark.Pass.ExtractKernels.DistributeNests-import Futhark.Pass.ExtractKernels.Distribution-import Futhark.Pass.ExtractKernels.ToGPU-import Futhark.Tools-import Futhark.Transform.FirstOrderTransform qualified as FOT-import Futhark.Util.Log-import Prelude hiding (log)---- | Convert the statements inside a map nest to kernel statements,--- attempting to parallelise any remaining (top-level) parallel--- statements.  Anything that is not a map, scan or reduction will--- simply be sequentialised.  This includes sequential loops that--- contain maps, scans or reduction.  In the future, we could probably--- do something more clever.  Make sure that the amount of parallelism--- to be exploited does not exceed the group size.  Further, as a hack--- we also consider the size of all intermediate arrays as--- "parallelism to be exploited" to avoid exploding shared memory.------ We distinguish between "minimum group size" and "maximum--- exploitable parallelism".-intrablockParallelise ::-  (MonadFreshNames m, LocalScope GPU m) =>-  KernelNest ->-  Lambda SOACS ->-  m-    ( Maybe-        ( (SubExp, SubExp),-          SubExp,-          Log,-          Stms GPU,-          Stms GPU-        )-    )-intrablockParallelise knest lam = runMaybeT $ do-  (ispace, inps) <- lift $ flatKernel knest--  (num_tblocks, w_stms) <--    lift $-      runBuilder $-        letSubExp "intra_num_tblocks"-          =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) (map snd ispace)--  let body = lambdaBody lam--  tblock_size <- newVName "computed_tblock_size"-  (wss_min, wss_avail, log, kbody) <--    lift . localScope (scopeOfLParams $ lambdaParams lam) $-      intrablockParalleliseBody body--  outside_scope <- lift askScope-  -- outside_scope may also contain the inputs, even though those are-  -- not actually available outside the kernel.-  let available v =-        v `M.member` outside_scope-          && v `notElem` map kernelInputName inps-  unless (allNames available $ freeIn (wss_min ++ wss_avail)) $-    fail "Irregular parallelism"--  ((intra_avail_par, kspace, read_input_stms), prelude_stms) <- lift $-    runBuilder $ do-      let foldBinOp' _ [] = eSubExp $ intConst Int64 1-          foldBinOp' bop (x : xs) = foldBinOp bop x xs-      ws_min <--        mapM (letSubExp "one_intra_par_min" <=< foldBinOp' (Mul Int64 OverflowUndef)) $-          filter (not . null) wss_min-      ws_avail <--        mapM (letSubExp "one_intra_par_avail" <=< foldBinOp' (Mul Int64 OverflowUndef)) $-          filter (not . null) wss_avail--      -- The amount of parallelism available *in the worst case* is-      -- equal to the smallest parallel loop, or *at least* 1.-      intra_avail_par <--        letSubExp "intra_avail_par" =<< foldBinOp' (SMin Int64) ws_avail--      -- The group size is either the maximum of the minimum parallelism-      -- exploited, or the desired parallelism (bounded by the max group-      -- size) in case there is no minimum.-      letBindNames [tblock_size]-        =<< if null ws_min-          then-            eBinOp-              (SMin Int64)-              (eSubExp =<< letSubExp "max_tblock_size" (Op $ SizeOp $ GetSizeMax SizeThreadBlock))-              (eSubExp intra_avail_par)-          else foldBinOp' (SMax Int64) ws_min--      let inputIsUsed input = kernelInputName input `nameIn` freeIn body-          used_inps = filter inputIsUsed inps--      addStms w_stms-      read_input_stms <- runBuilder_ $ mapM readGroupKernelInput used_inps-      space <- SegSpace <$> newVName "phys_tblock_id" <*> pure ispace-      pure (intra_avail_par, space, read_input_stms)--  let kbody' = kbody {bodyStms = read_input_stms <> bodyStms kbody}--  let nested_pat = loopNestingPat first_nest-      rts = map (length ispace `stripArray`) $ patTypes nested_pat-      grid = KernelGrid (Count num_tblocks) (Count $ Var tblock_size)-      lvl = SegBlock SegNoVirt (Just grid)-      kstm = Let nested_pat aux $ Op $ SegOp $ SegMap lvl kspace rts kbody'--  let intra_min_par = intra_avail_par-  pure-    ( (intra_min_par, intra_avail_par),-      Var tblock_size,-      log,-      prelude_stms,-      oneStm kstm-    )-  where-    first_nest = fst knest-    aux = loopNestingAux first_nest--readGroupKernelInput ::-  (DistRep (Rep m), MonadBuilder m) =>-  KernelInput ->-  m ()-readGroupKernelInput inp-  | Array {} <- kernelInputType inp = do-      v <- newName $ kernelInputName inp-      readKernelInput inp {kernelInputName = v}-      letBindNames [kernelInputName inp] $ BasicOp $ Replicate mempty $ Var v-  | otherwise =-      readKernelInput inp--data IntraAcc = IntraAcc-  { accMinPar :: S.Set [SubExp],-    accAvailPar :: S.Set [SubExp],-    accLog :: Log-  }--instance Semigroup IntraAcc where-  IntraAcc min_x avail_x log_x <> IntraAcc min_y avail_y log_y =-    IntraAcc (min_x <> min_y) (avail_x <> avail_y) (log_x <> log_y)--instance Monoid IntraAcc where-  mempty = IntraAcc mempty mempty mempty--type IntrablockM =-  BuilderT GPU (RWS () IntraAcc VNameSource)--instance MonadLogger IntrablockM where-  addLog log = tell mempty {accLog = log}--runIntrablockM ::-  (MonadFreshNames m, HasScope GPU m) =>-  IntrablockM () ->-  m (IntraAcc, Stms GPU)-runIntrablockM m = do-  scope <- castScope <$> askScope-  modifyNameSource $ \src ->-    let (((), kstms), src', acc) = runRWS (runBuilderT m scope) () src-     in ((acc, kstms), src')--parallelMin :: [SubExp] -> IntrablockM ()-parallelMin ws =-  tell-    mempty-      { accMinPar = S.singleton ws,-        accAvailPar = S.singleton ws-      }--intrablockBody :: Body SOACS -> IntrablockM (Body GPU)-intrablockBody body = do-  stms <- collectStms_ $ intrablockStms $ bodyStms body-  pure $ mkBody stms $ bodyResult body--intrablockLambda :: Lambda SOACS -> IntrablockM (Lambda GPU)-intrablockLambda lam =-  mkLambda (lambdaParams lam) $-    bodyBind =<< intrablockBody (lambdaBody lam)--intrablockWithAccInput :: WithAccInput SOACS -> IntrablockM (WithAccInput GPU)-intrablockWithAccInput (shape, arrs, Nothing) =-  pure (shape, arrs, Nothing)-intrablockWithAccInput (shape, arrs, Just (lam, nes)) = do-  lam' <- intrablockLambda lam-  pure (shape, arrs, Just (lam', nes))--intrablockStm :: Stm SOACS -> IntrablockM ()-intrablockStm stm@(Let pat aux e) = do-  scope <- askScope-  let lvl = SegThreadInBlock SegNoVirt--  case e of-    Loop merge form loopbody ->-      localScope (scopeOfLoopForm form <> scopeOfFParams (map fst merge)) $ do-        loopbody' <- intrablockBody loopbody-        certifying (stmAuxCerts aux) . letBind pat $-          Loop merge form loopbody'-    Match cond cases defbody ifdec -> do-      cases' <- mapM (traverse intrablockBody) cases-      defbody' <- intrablockBody defbody-      certifying (stmAuxCerts aux) . letBind pat $-        Match cond cases' defbody' ifdec-    WithAcc inputs lam -> do-      inputs' <- mapM intrablockWithAccInput inputs-      lam' <- intrablockLambda lam-      certifying (stmAuxCerts aux) . letBind pat $ WithAcc inputs' lam'-    Op soac-      | "sequential_outer" `inAttrs` stmAuxAttrs aux ->-          intrablockStms . fmap (certify (stmAuxCerts aux))-            =<< runBuilder_ (FOT.transformSOAC pat soac)-    Op (Screma w arrs form)-      | Just lam <- isMapSOAC form -> do-          let loopnest = MapNesting pat aux w $ zip (lambdaParams lam) arrs-              env =-                DistEnv-                  { distNest =-                      singleNesting $ Nesting mempty loopnest,-                    distScope =-                      scopeOfPat pat-                        <> scopeForGPU (scopeOf lam)-                        <> scope,-                    distOnInnerMap =-                      distributeMap,-                    distOnTopLevelStms =-                      liftInner . collectStms_ . intrablockStms,-                    distSegLevel = \minw _ _ -> do-                      lift $ parallelMin minw-                      pure lvl,-                    distOnSOACSStms =-                      pure . oneStm . soacsStmToGPU,-                    distOnSOACSLambda =-                      pure . soacsLambdaToGPU-                  }-              acc =-                DistAcc-                  { distTargets = singleTarget (pat, bodyResult $ lambdaBody lam),-                    distStms = mempty-                  }--          addStms-            =<< runDistNestT env (distributeMapBodyStms acc (bodyStms $ lambdaBody lam))-    Op (Screma w arrs form)-      | Just (post_lam, scans, mapfun) <- isMaposcanomapSOAC form,-        -- FIXME: Futhark.CodeGen.ImpGen.GPU.Block.compileGroupOp-        -- cannot handle multiple scan operators yet.-        Scan scanfun nes <- singleScan scans -> do-          let scanfun' = soacsLambdaToGPU scanfun-              mapfun' = soacsLambdaToGPU mapfun-              post_op = SegPostOp $ soacsLambdaToGPU post_lam-          certifying (stmAuxCerts aux) $-            addStms =<< segScan lvl pat mempty w [SegBinOp Noncommutative scanfun' nes mempty] mapfun' post_op arrs [] []-          parallelMin [w]-    Op (Screma w arrs form)-      | Just (reds, map_lam) <- isRedomapSOAC form -> do-          let onReduce (Reduce comm red_lam nes) =-                SegBinOp comm (soacsLambdaToGPU red_lam) nes mempty-              reds' = map onReduce reds-              map_lam' = soacsLambdaToGPU map_lam-          certifying (stmAuxCerts aux) $-            addStms =<< segRed lvl pat mempty w reds' map_lam' arrs [] []-          parallelMin [w]-    Op (Screma w arrs form) ->-      -- This screma is too complicated for us to immediately do-      -- anything, so split it up and try again.-      mapM_ intrablockStm . fmap (certify (stmAuxCerts aux)) . snd-        =<< runBuilderT (dissectScrema pat w form arrs) (scopeForSOACs scope)-    Op (Hist w arrs ops bucket_fun) -> do-      ops' <- forM ops $ \(HistOp num_bins rf dests nes op) -> do-        (op', nes', shape) <- determineReduceOp op nes-        let op'' = soacsLambdaToGPU op'-        pure $ GPU.HistOp num_bins rf dests nes' shape op''--      let bucket_fun' = soacsLambdaToGPU bucket_fun-      certifying (stmAuxCerts aux) $-        addStms =<< segHist lvl pat w [] [] ops' bucket_fun' arrs-      parallelMin [w]-    Op (Stream w arrs accs lam)-      | chunk_size_param : _ <- lambdaParams lam -> do-          types <- asksScope castScope-          ((), stream_stms) <--            runBuilderT (sequentialStreamWholeArray pat w accs lam arrs) types-          let replace (Var v) | v == paramName chunk_size_param = w-              replace se = se-              replaceSets (IntraAcc x y log) =-                IntraAcc (S.map (map replace) x) (S.map (map replace) y) log-          censor replaceSets $ intrablockStms stream_stms-    _ ->-      addStm $ soacsStmToGPU stm--intrablockStms :: Stms SOACS -> IntrablockM ()-intrablockStms = mapM_ intrablockStm--intrablockParalleliseBody ::-  (MonadFreshNames m, HasScope GPU m) =>-  Body SOACS ->-  m ([[SubExp]], [[SubExp]], Log, KernelBody GPU)-intrablockParalleliseBody body = do-  (IntraAcc min_ws avail_ws log, kstms) <--    runIntrablockM $ intrablockStms $ bodyStms body-  pure-    ( S.toList min_ws,-      S.toList avail_ws,-      log,-      Body () kstms $ map ret $ bodyResult body-    )-  where-    ret (SubExpRes cs se) = Returns ResultMaySimplify cs se
− src/Futhark/Pass/ExtractKernels/StreamKernel.hs 
@@ -1,82 +0,0 @@-{-# LANGUAGE TypeFamilies #-}--module Futhark.Pass.ExtractKernels.StreamKernel-  ( segThreadCapped,-  )-where--import Control.Monad-import Data.List ()-import Futhark.Analysis.PrimExp-import Futhark.IR-import Futhark.IR.GPU hiding-  ( BasicOp,-    Body,-    Exp,-    FParam,-    FunDef,-    LParam,-    Lambda,-    Pat,-    PatElem,-    Prog,-    RetType,-    Stm,-  )-import Futhark.MonadFreshNames-import Futhark.Pass.ExtractKernels.BlockedKernel-import Futhark.Pass.ExtractKernels.ToGPU-import Futhark.Tools-import Prelude hiding (quot)--data KernelSize = KernelSize-  { -- | Int64-    kernelElementsPerThread :: SubExp,-    -- | Int32-    kernelNumThreads :: SubExp-  }-  deriving (Eq, Ord, Show)--numberOfBlocks ::-  (MonadBuilder m, Op (Rep m) ~ HostOp inner (Rep m)) =>-  Name ->-  SubExp ->-  SubExp ->-  m (SubExp, SubExp)-numberOfBlocks desc w tblock_size = do-  max_num_tblocks_key <- nameFromText . prettyText <$> newVName (desc <> "_num_tblocks")-  num_tblocks <--    letSubExp "num_tblocks" $-      Op $-        SizeOp $-          CalcNumBlocks w max_num_tblocks_key tblock_size-  num_threads <--    letSubExp "num_threads" $-      BasicOp $-        BinOp (Mul Int64 OverflowUndef) num_tblocks tblock_size-  pure (num_tblocks, num_threads)---- | Like 'segThread', but cap the thread count to the input size.--- This is more efficient for small kernels, e.g. summing a small--- array.-segThreadCapped :: (MonadFreshNames m) => MkSegLevel GPU m-segThreadCapped ws desc r = do-  w <--    letSubExp "nest_size"-      =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) ws-  tblock_size <- getSize (desc <> "_tblock_size") SizeThreadBlock--  case r of-    ManyThreads -> do-      usable_groups <--        letSubExp "segmap_usable_groups"-          =<< eBinOp-            (SDivUp Int64 Unsafe)-            (eSubExp w)-            (eSubExp =<< asIntS Int64 tblock_size)-      let grid = KernelGrid (Count usable_groups) (Count tblock_size)-      pure $ SegThread SegNoVirt (Just grid)-    NoRecommendation v -> do-      (num_tblocks, _) <- numberOfBlocks desc w tblock_size-      let grid = KernelGrid (Count num_tblocks) (Count tblock_size)-      pure $ SegThread v (Just grid)
− src/Futhark/Pass/ExtractKernels/ToGPU.hs 
@@ -1,84 +0,0 @@-{-# LANGUAGE TypeFamilies #-}--module Futhark.Pass.ExtractKernels.ToGPU-  ( getSize,-    segThread,-    soacsLambdaToGPU,-    soacsStmToGPU,-    scopeForGPU,-    scopeForSOACs,-    injectSOACS,-  )-where--import Control.Monad.Identity-import Data.List ()-import Futhark.IR-import Futhark.IR.GPU-import Futhark.IR.SOACS (SOACS)-import Futhark.IR.SOACS.SOAC qualified as SOAC-import Futhark.Tools--getSize ::-  (MonadBuilder m, Op (Rep m) ~ HostOp inner (Rep m)) =>-  Name ->-  SizeClass ->-  m SubExp-getSize desc size_class = do-  size_key <- nameFromText . prettyText <$> newVName desc-  letSubExp desc $ Op $ SizeOp $ GetSize size_key size_class--segThread ::-  (MonadBuilder m, Op (Rep m) ~ HostOp inner (Rep m)) =>-  Name ->-  m SegLevel-segThread desc =-  SegThread SegVirt <$> (Just <$> kernelGrid)-  where-    kernelGrid =-      KernelGrid-        <$> (Count <$> getSize (desc <> "_num_tblocks") SizeGrid)-        <*> (Count <$> getSize (desc <> "_tblock_size") SizeThreadBlock)--injectSOACS ::-  ( Monad m,-    SameScope from to,-    ExpDec from ~ ExpDec to,-    BodyDec from ~ BodyDec to,-    RetType from ~ RetType to,-    BranchType from ~ BranchType to,-    Op from ~ SOAC from-  ) =>-  (SOAC to -> Op to) ->-  Rephraser m from to-injectSOACS f =-  Rephraser-    { rephraseExpDec = pure,-      rephraseBodyDec = pure,-      rephraseLetBoundDec = pure,-      rephraseFParamDec = pure,-      rephraseLParamDec = pure,-      rephraseOp = fmap f . onSOAC,-      rephraseRetType = pure,-      rephraseBranchType = pure-    }-  where-    onSOAC = SOAC.mapSOACM mapper-    mapper =-      SOAC.SOACMapper-        { SOAC.mapOnSOACSubExp = pure,-          SOAC.mapOnSOACVName = pure,-          SOAC.mapOnSOACLambda = rephraseLambda $ injectSOACS f-        }--soacsStmToGPU :: Stm SOACS -> Stm GPU-soacsStmToGPU = runIdentity . rephraseStm (injectSOACS OtherOp)--soacsLambdaToGPU :: Lambda SOACS -> Lambda GPU-soacsLambdaToGPU = runIdentity . rephraseLambda (injectSOACS OtherOp)--scopeForSOACs :: Scope GPU -> Scope SOACS-scopeForSOACs = castScope--scopeForGPU :: Scope SOACS -> Scope GPU-scopeForGPU = castScope
src/Futhark/Pass/ExtractMulticore.hs view
@@ -2,7 +2,7 @@  -- | Extraction of parallelism from a SOACs program.  This generates -- parallel constructs aimed at CPU execution, which in particular may--- involve ad-hoc irregular nested parallelism.+-- involve ad-hoc nonuniform nested parallelism. module Futhark.Pass.ExtractMulticore (extractMulticore) where  import Control.Monad@@ -23,10 +23,12 @@   ) import Futhark.IR.SOACS qualified as SOACS import Futhark.Pass-import Futhark.Pass.ExtractKernels.DistributeNests-import Futhark.Pass.ExtractKernels.ToGPU (injectSOACS)+import Futhark.Pass.Flatten.Builtins (determineReduceOp)+import Futhark.Pass.Flatten.Incremental (lambdaHasParallelism) import Futhark.Tools+import Futhark.Transform.FirstOrderTransform qualified as FOT import Futhark.Transform.Rename (Rename, renameSomething)+import Futhark.Transform.ToGPU (injectSOACS) import Futhark.Util.Log  newtype ExtractM a = ExtractM (ReaderT (Scope MC) (State VNameSource) a)@@ -215,10 +217,19 @@   error "transformSOAC: unhandled VJP" transformSOAC _ _ WithVJP {} =   error "transformSOAC: unhandled WithVJP"+transformSOAC pat _ (FlatMap w arrs lam) = do+  -- Sequentialise the FlatMap itself (but not its contents) via the first-order+  -- transform, then transform the resulting stms. This does lose us+  -- parallelism, but hopefully it is not often the case that the FlatMap is the+  -- only source of parallelism.+  soacs_scope <- castScope <$> askScope+  flatmap_stms <-+    flip runBuilderT_ soacs_scope $ FOT.transformFlatMap pat w arrs lam+  transformStms flatmap_stms transformSOAC pat _ (Screma w arrs form)   | Just lam <- isMapSOAC form = do       seq_op <- transformMap DoNotRename sequentialiseBody w lam arrs-      if lambdaContainsParallelism lam+      if lambdaHasParallelism (const False) lam         then do           par_op <- transformMap DoRename transformBody w lam arrs           pure $ oneStm (Let pat (defAux ()) $ Op $ ParOp (Just par_op) seq_op)@@ -226,7 +237,7 @@   | Just (reds, map_lam) <- isRedomapSOAC form = do       (seq_reds_stms, seq_op) <-         transformRedomap DoNotRename sequentialiseBody w reds map_lam arrs-      if lambdaContainsParallelism map_lam+      if lambdaHasParallelism (const False) map_lam         then do           (par_reds_stms, par_op) <-             transformRedomap DoRename transformBody w reds map_lam arrs@@ -259,7 +270,7 @@   (seq_hist_stms, seq_op) <-     transformHist DoNotRename sequentialiseBody w hists map_lam arrs -  if lambdaContainsParallelism map_lam+  if lambdaHasParallelism (const False) map_lam     then do       (par_hist_stms, par_op) <-         transformHist DoRename transformBody w hists map_lam arrs
+ src/Futhark/Pass/Flatten.hs view
@@ -0,0 +1,715 @@+{-# LANGUAGE TypeFamilies #-}++-- | This pass transforms parallelism expressed with arbitrarily nested SOACs to+-- instead be expressed with limited-nesting SegOps. This is the so-called+-- "flattening transformation" (sometimes called "vectorization", although we do+-- not use that term much in the Futhark compiler).+--+-- This is a sophisticated pass that does various clever things:+--+-- - Detects uniform nesting and flattens it more efficiently than the+--   nonuniform case.+--+-- - DPH-style vectorization avoidance.+--+-- - Incremental flattening ("Futhark.Pass.Flatten.Incremental").+--+-- - Intrablock flattening ("Futhark.Pass.Flatten.Intrablock").+--+-- The goal is that *any* Futhark program must be compilable parallel GPU code,+-- although in some cases the resulting code is not particularly efficient.+--+-- The idea is to perform distribution on one level at a time, and produce+-- "irregular maps" that can accept and produce irregular arrays. These+-- irregular maps will then be transformed into flat parallelism based on their+-- contents. If irregular maps contain only a single Stm, then it is fairly+-- straightforward, as we simply implement flattening rules for every single+-- kind of expression. Of course that is also somewhat inefficient, so we want+-- to support multiple Stms for things like scalar code.+--+-- Nomenclature:+--+-- A /map-nest/ is the collection of parallel operations enclosing some code. For+-- simplicity, we say "map-nest" even when the top level parallel operation is+-- actually a redomap or other screma.+--+-- An /irregular array/ is a multidimensional array like '[[1,2],[3]]', where rows+-- have different shapes. These are not directly supported in Futhark or in the+-- Futhark IR, but are encoded in various ways.+--+-- We say that an operation or type in a map-nest is /uniform/ when its size+-- (including internal sizes and sizes of inputs) and control flow is invariant+-- to the map-nest. Converse, it is /nonuniform/ when it is variant. When we+-- distribute a uniform statement, the intermediate results are regular, and+-- otherwise irregular. A statement that uses an irregular array is necessarily+-- nonuniform.+--+-- Take care not to confuse the terms "regular" and "uniform" - we say "regular"+-- only about arrays! "Uniform" is the general concept.+--+-- /Uniform nested parallelism/ is nested parallelism whose size is uniform to+-- the enclosing map nest, and which uses only variables whose types are+-- uniform, and which is enclosed in uniform control flow. /Nonuniform nested+-- parallelism/ is the converse. Many of the optimisations here are about+-- detecting the uniform case. We previously often used the terms "regular+-- nested parallelism" and "irregular nested parallelism", but this is now+-- discouraged, as explained above.+module Futhark.Pass.Flatten (flattenSOACs) where++import Control.Monad+import Data.Bifunctor (second)+import Data.Foldable+import Data.List qualified as L+import Data.Map qualified as M+import Data.Set qualified as S+import Futhark.Analysis.Alias (analyseBody)+import Futhark.IR.Aliases (Aliases, bodyAliases)+import Futhark.IR.GPU+import Futhark.IR.SOACS+import Futhark.MonadFreshNames+import Futhark.Pass+import Futhark.Pass.Flatten.BasicOp+import Futhark.Pass.Flatten.Builtins+import Futhark.Pass.Flatten.Distribute+import Futhark.Pass.Flatten.General+import Futhark.Pass.Flatten.Incremental+import Futhark.Pass.Flatten.Loop+import Futhark.Pass.Flatten.Match+import Futhark.Pass.Flatten.PreProcess+import Futhark.Pass.Flatten.SOAC+import Futhark.Pass.Flatten.WithAcc+import Futhark.Tools+import Futhark.Transform.FirstOrderTransform qualified as FOT+import Futhark.Transform.Rename+import Futhark.Transform.ToGPU (soacsLambdaToGPU, soacsStmToGPU)+import Prelude hiding (div, quot, rem)++type FunSizeParams = Name -> S.Set Int++-- | The irregularity handling mode requested by a statement, defaulting to the+-- mode already in effect. @#[flattening(sequentialise_nonuniform)]@ asks that+-- nonuniform nested parallelism be sequentialised rather than flattened; see+-- 'SequentialiseIrregularAll'.+irregularityFor :: DistIrregularity -> StmAux a -> DistIrregularity+irregularityFor irreg aux+  | AttrComp "flattening" ["sequentialise_nonuniform"] `inAttrs` stmAuxAttrs aux =+      SequentialiseIrregularAll+  | otherwise = irreg++flattenOpsFor :: FunHasParallelism -> FunSizeParams -> DistIrregularity -> SegLevel -> FlattenOps+flattenOpsFor funHasParallelism funSizeParams irreg lvl =+  FlattenOps+    { flattenSegLevel = lvl,+      flattenIrregularity = irreg,+      flattenFunHasParallelism = funHasParallelism,+      flattenDistStmWith = transformDistStm funSizeParams,+      flattenScalarStmAt = transformScalarStm,+      flattenTopLevelStm = transformStm funHasParallelism funSizeParams+    }++transformScalarStms ::+  SegLevel ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  Stms SOACS ->+  FlattenM DistEnv+transformScalarStms lvl segments env inps distres stms = do+  let bound_in_batch = namesFromList $ concatMap (patNames . stmPat) $ stmsToList stms+      allCerts = foldMap (\stm -> distCerts inps (stmAux stm) env) (stmsToList stms)+      certs = Certs $ filter (`notNameIn` bound_in_batch) $ unCerts allCerts+  vs <- certifying certs $ letTupExp "scalar_dist" <=< renameExp <=< segMap lvl segments $ \is -> do+    readInputs segments env (toList is) inps+    addStms $ fmap soacsStmToGPU stms+    pure $ subExpsRes $ map (Var . distResName) distres+  insertRepsM (zip (map distResTag distres) $ map Regular vs) env++transformScalarStm ::+  SegLevel ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  Stm SOACS ->+  FlattenM DistEnv+transformScalarStm lvl segments env inps res stm =+  transformScalarStms lvl segments env inps res (oneStm stm)++-- | Transform a top-level 'Screma' by treating it as the empty-'Segments' case+-- of a nested one: it is not enclosed in any map-nest, so there are no+-- segments, the mapped arrays are plain regular top-level values+-- ('DistInputFree'), and the results are necessarily regular.+transformTopLevelScrema ::+  FunHasParallelism ->+  FunSizeParams ->+  Pat Type ->+  StmAux () ->+  SubExp ->+  [VName] ->+  ScremaForm SOACS ->+  FlattenM ()+transformTopLevelScrema funHasParallelism funSizeParams pat aux w arrs form = do+  let irreg = irregularityFor DistributeIrregular aux+      ops = flattenOpsFor funHasParallelism funSizeParams irreg defaultSegLevel+  arr_ts <- mapM lookupType arrs+  -- 'flattenScrema' may bind the names of the pattern it is given (some paths+  -- bind them directly, others only insert reps), so we pass it a fresh pattern+  -- and bind the real pattern names ourselves from the result.+  nested_pat <- renamePat pat+  let inps = zipWith (\arr t -> (arr, DistInputFree arr t)) arrs arr_ts+      res =+        zipWith+          (\i pe -> DistResult (ResTag i) (DistType [] (Rank 0) (patElemType pe)) (patElemName pe))+          [0 ..]+          (patElems nested_pat)+  env <- flattenScrema ops [] (DistEnv mempty) inps res (nested_pat, aux) (w, arrs, form)+  forM_ (zip (patNames pat) res) $ \(pat_v, r) ->+    case resVar (distResTag r) env of+      Regular v ->+        letBindNames [pat_v] $ BasicOp $ SubExp $ Var v+      Irregular _ ->+        error "transformTopLevelScrema: top-level result cannot be irregular"++liftArg :: SegLevel -> Segments -> SubExp -> DistInputs -> DistEnv -> (SubExp, Diet) -> FlattenM [(SubExp, Diet)]+liftArg lvl segments w inps env (se, d) = do+  (_, rep) <- liftSubExp lvl segments inps env se+  case rep of+    Regular v -> do+      v_t <- lookupType v+      v' <-+        if arrayShape v_t == Shape [w]+          then pure v+          else+            letExp "lifted_arg_flat" . BasicOp $+              Reshape v $+                reshapeAll (arrayShape v_t) (Shape [w])+      pure [(Var v', d)]+    Irregular irreg -> do+      vs <- irregularRepToFlatArrs w irreg+      -- Only apply the original diet to the 'elems' array.+      pure $ zip (map Var vs) $ replicate 4 Observe ++ [d]++liftRegArg :: SegLevel -> Segments -> SubExp -> DistInputs -> DistEnv -> (SubExp, Diet) -> FlattenM (SubExp, Diet)+liftRegArg lvl _segments w inps env (se, d) = do+  se_t <- subExpInputType inps se+  let se_shape = arrayShape se_t+      expected_shape = Shape [w] <> se_shape+  v <- liftSubExpRegular lvl [w] inps env expected_shape se+  pure (Var v, d)++-- Lifts a functions return type such that it matches the lifted functions+-- return type.+--+-- A lifted function corresponds to 'map f', which always produces fresh arrays.+-- We therefore mark all array components of the return type as 'Unique', such+-- that the results are known to not alias anything (in particular not the+-- arguments). Maintaining this invariant may require inserting copies in the+-- function body; see 'freshenResult'.+liftRetType :: SubExp -> [RetType SOACS] -> [RetType GPU]+liftRetType w = concat . snd . L.mapAccumL liftType 0+  where+    liftType i rettype =+      let lifted = case rettype of+            Prim pt -> pure $ arrayOf (Prim pt) (Shape [Free w]) Unique+            Array pt _ _ ->+              let num_data = Prim int64+                  segs = arrayOf (Prim int64) (Shape [Free w]) Unique+                  flags = arrayOf (Prim Bool) (Shape [Ext i]) Unique+                  offsets = arrayOf (Prim int64) (Shape [Free w]) Unique+                  elems = arrayOf (Prim pt) (Shape [Ext i]) Unique+               in [num_data, segs, flags, offsets, elems]+            Acc {} -> error "liftRetType: Acc"+            Mem {} -> error "liftRetType: Mem"+       in (i + length lifted, lifted)++liftRegularRetType :: DistInputs -> SubExp -> [RetType SOACS] -> [RetType GPU]+liftRegularRetType inps w = concat . snd . L.mapAccumL liftType 0+  where+    liftType i rettype =+      let lifted = case rettype of+            Prim pt -> pure $ arrayOf (Prim pt) (Shape [Free w]) Unique+            Array pt shape _ ->+              if needsIrregularRetType inps rettype+                then+                  let num_data = Prim int64+                      segs = arrayOf (Prim int64) (Shape [Free w]) Unique+                      flags = arrayOf (Prim Bool) (Shape [Ext i]) Unique+                      offsets = arrayOf (Prim int64) (Shape [Free w]) Unique+                      elems = arrayOf (Prim pt) (Shape [Ext i]) Unique+                   in [num_data, segs, flags, offsets, elems]+                else+                  pure $ arrayOf (Prim pt) (Shape [Free w] <> shape) Unique+            Acc {} -> error "liftRetType: Acc"+            Mem {} -> error "liftRetType: Mem"+       in (i + length lifted, lifted)++liftFunName :: Name -> Name+liftFunName name = name <> "_lifted"++liftUniformFunName :: Name -> Name+liftUniformFunName name = name <> "_uniform_lifted"++flattenApply ::+  FunSizeParams ->+  SegLevel ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  (Pat Type, StmAux ()) ->+  (Name, [(SubExp, Diet)], [(RetType SOACS, RetAls)], Safety) ->+  FlattenM DistEnv+flattenApply funSizeParams lvl segments env inps res (pat, aux) (name, args, rettype, s) =+  case lvl of+    SegThread {} -> do+      let size_positions = funSizeParams name+          indexed_args = zip [0 ..] args+          isSizeArg = (`S.member` size_positions) . fst+          (size_args, value_args) = L.partition isSizeArg indexed_args+      let nonuniform = any (isVariant inps . fst . snd) size_args+          name' = if nonuniform then liftFunName name else liftUniformFunName name+          mode = if nonuniform then NonUniformLift else UniformLift+      demandLifted name mode+      w <- letSubExp "num_segments" =<< toExp (segmentCount segments)++      args' <-+        if nonuniform+          then+            ((w, Observe) :) . concat <$> mapM (liftArg lvl segments w inps env) args+          else do+            value_args' <- mapM (liftRegArg lvl segments w inps env . snd) value_args+            -- We do not lift 'size_args' because they correspond to size+            -- parameters, which are invariant in the uniform case.+            pure $ (w, Observe) : map snd size_args <> value_args'+      args_ts <- mapM (subExpType . fst) args'+      let dietToUnique Consume = Unique+          dietToUnique Observe = Nonunique+          param_ts = zipWith toDecl args_ts $ map (dietToUnique . snd) args'+          rettype' =+            if nonuniform+              then addRetAls param_ts $ liftRetType w $ map fst rettype+              else addRetAls param_ts $ liftRegularRetType inps w $ map fst rettype+      result <- letTupExp (name' <> "_res") $ Apply name' args' rettype' s+      let reps =+            if nonuniform+              then resultToResReps (map fst rettype) result+              -- XXX: This could instead distinguish between regular and+              -- irregular results based on their return types.+              else resultToResRepsByDistResult res result+      reps' <- zipWithM (reshapeLiftedApplyResult segments) (map fst rettype) reps+      insertRepsM (zip (map distResTag res) reps') env+    -- TODO: we currently do not handle intrablock function applications. It+    -- is possible we could do intrablock-level lifting of functions, but+    -- for now, we simply do not generate intrablock kernels if they would+    -- contain calls to parallel functions.+    _ ->+      if all isRegularDistResult res+        then transformScalarStm lvl segments env inps res $ Let pat aux (Apply name args rettype s)+        else error "Unhandled Apply in non SegThread Seglevel"++transformDistStm :: FunSizeParams -> FlattenOps -> Segments -> DistEnv -> DistStm -> FlattenM DistEnv+transformDistStm _ outer_ops segments env (DistStm inps res (ScalarStm stms)) =+  transformScalarStms (flattenSegLevel outer_ops) segments env inps res stms+transformDistStm funSizeParams outer_ops segments env (DistStm inps res (ParallelStm stm)) = do+  case stm of+    Let pat aux (BasicOp e) -> do+      let ~[res'] = res+          ~[pe] = patElems pat+      flattenBasicOp ops segments env (inps, res', pe, aux, e)+    Let pat aux (Op (Screma w arrs form)) ->+      flattenScrema ops segments env inps res (pat, aux) (w, arrs, form)+    Let _ aux (Match scrutinees cases defaultCase rt) ->+      flattenMatch ops segments env inps res aux scrutinees cases defaultCase rt+    Let pat aux (Apply name args rettype s) ->+      flattenApply funSizeParams lvl segments env inps res (pat, aux) (name, args, rettype, s)+    Let pat aux (Loop merge (ForLoop i it n) body) ->+      flattenLoop ops segments env inps res (pat, aux) (merge, ForLoop i it n, body)+    Let pat aux (Loop merge (WhileLoop cond) body) -> do+      flattenLoop ops segments env inps res (pat, aux) (merge, WhileLoop cond, body)+    Let pat aux (WithAcc inputs lam) ->+      flattenWithAcc ops segments env inps res pat aux inputs lam+    (Let pat aux (Op (Hist w hist_inputs hist_ops bucket_fun))) ->+      flattenHist ops segments env inps res (pat, aux) (w, hist_inputs, hist_ops, bucket_fun)+    Let _ aux (Op (FlatMap w arrs lam)) ->+      flattenFlatMapNested ops segments env inps res aux w arrs lam+    Let _ _ (Op (Stream {})) -> error "transformDistStm: Stream should have been removed"+    Let _ _ (Op (JVP {})) -> error "Unhandled JVP"+    Let _ _ (Op (VJP {})) -> error "Unhandled VJP"+    Let _ _ (Op (WithVJP {})) -> error "Unhandled WithVJP"+  where+    lvl = flattenSegLevel outer_ops+    ops =+      outer_ops+        { flattenIrregularity =+            irregularityFor (flattenIrregularity outer_ops) (stmAux stm)+        }++reshapeLiftedApplyResult :: Segments -> RetType SOACS -> ResRep -> FlattenM ResRep+reshapeLiftedApplyResult segments Prim {} (Regular v) = do+  v_t <- lookupType v+  let expectedShape = segmentsShape segments+  v' <-+    if arrayShape v_t == expectedShape+      then pure v+      else+        letExp "lifted_apply_res" . BasicOp $+          Reshape v $+            reshapeAll (arrayShape v_t) expectedShape+  pure $ Regular v'+reshapeLiftedApplyResult _ _ rep =+  pure rep++liftBody :: FunHasParallelism -> FunSizeParams -> SegLevel -> SubExp -> DistInputs -> DistEnv -> DistStms -> Result -> FlattenM Result+liftBody funHasParallelism funSizeParams lvl w inputs env dstms result = do+  let segments = [w]+      ops = flattenOpsFor funHasParallelism funSizeParams DistributeIrregular lvl+  env' <- foldM (flattenDistStm ops segments) env dstms+  result' <- mapM (liftResult lvl segments inputs env') result+  pure $ concat result'++liftUniformFunBody :: FunHasParallelism -> FunSizeParams -> SegLevel -> SubExp -> DistInputs -> DistEnv -> DistStms -> [RetType SOACS] -> Result -> FlattenM Result+liftUniformFunBody funHasParallelism funSizeParams lvl w inputs env dstms rettype result = do+  let segments = [w]+      ops = flattenOpsFor funHasParallelism funSizeParams DistributeIrregular lvl+  env' <- foldM (flattenDistStm ops segments) env dstms+  concat <$> zipWithM (liftRegResult lvl segments w inputs env') rettype result++-- | A lifted function must return fresh, non-aliasing arrays (as it+-- corresponds to 'map f'; see 'liftRetType').  This is not+-- automatically the case: a result may alias a parameter (when a value+-- is passed straight through), or the same array may be returned in+-- multiple result positions (which happens for functions that return+-- the same value more than once).  For every such result we insert a+-- copy to re-establish the invariant.  Results that are already fresh+-- are left untouched, so no superfluous copies are inserted.+freshenResult :: [FParam GPU] -> FlattenM Result -> FlattenM Result+freshenResult params m = do+  (result, stms) <- collectStms m+  addStms stms+  let param_names = namesFromList $ map paramName params+      -- Transitive aliases of each result, including aliases with+      -- parameters and other results.+      als = bodyAliases (analyseBody mempty (Body () stms result) :: Body (Aliases GPU))+  reverse . snd <$> foldM freshen (param_names, []) (zip result als)+  where+    freshen (taken, acc) (SubExpRes cs (Var v), v_als) = do+      v_t <- lookupType v+      case v_t of+        Array {}+          | taken `namesIntersect` v_als -> do+              v' <- letExp "fresh_result" $ BasicOp $ Replicate mempty $ Var v+              pure (taken, SubExpRes cs (Var v') : acc)+        _ ->+          pure (taken <> v_als, SubExpRes cs (Var v) : acc)+    freshen (taken, acc) (res', _) =+      pure (taken, res' : acc)++analyseFunParallelism :: [FunDef SOACS] -> M.Map Name Bool+analyseFunParallelism funs =+  M.fromList [(funDefName fun, hasParallelFun mempty (funDefName fun)) | fun <- funs]+  where+    funsByName =+      M.fromList [(funDefName fun, fun) | fun <- funs]+    hasParallelFun seen fname+      | isBuiltInFunction fname =+          False+      -- avoid cycles even thought it is impossible now+      | fname `S.member` seen =+          False+      | Just fun <- M.lookup fname funsByName =+          any (isParallelStm (hasParallelFun (S.insert fname seen))) $+            bodyStms $+              funDefBody fun+      | otherwise =+          error $ "analyseFunParallelism: unknown function " ++ prettyString fname++analyseFunSizeParams :: [FunDef SOACS] -> M.Map Name (S.Set Int)+analyseFunSizeParams = M.fromList . map analyse+  where+    analyse fd =+      let fparams = funDefParams fd+          rettype = funDefRetType fd+          size_names = freeIn (map paramType fparams, map fst rettype)+          isSizeParam p = paramName p `nameIn` size_names+          indexed_params = zip [0 ..] fparams+          size_params = filter (isSizeParam . snd) indexed_params+       in (funDefName fd, S.fromList $ map fst size_params)++addRetAls :: [DeclType] -> [RetType GPU] -> [(RetType GPU, RetAls)]+addRetAls params rettype = zip rettype $ map possibleAliases rettype+  where+    aliasable (Array _ _ Nonunique) = True+    aliasable _ = False+    aliasable_params =+      map snd $ filter (aliasable . fst) $ zip params [0 ..]+    aliasable_rets =+      map snd $ filter (aliasable . declExtTypeOf . fst) $ zip rettype [0 ..]+    possibleAliases t+      | aliasable t = RetAls aliasable_params aliasable_rets+      | otherwise = mempty++liftFunDef ::+  FunHasParallelism ->+  FunSizeParams ->+  Scope SOACS ->+  FunDef SOACS ->+  PassM (FunDef GPU, S.Set DemandFn)+liftFunDef funHasParallelism funSizeParams const_scope fd = do+  let FunDef+        { funDefBody = body,+          funDefParams = fparams,+          funDefRetType = rettype+        } = fd+  wp <- newParam "w" $ Prim int64+  let w = Var $ paramName wp+  (fparams', reps) <- mapAndUnzipM (liftParam w) fparams+  let fparams'' = wp : concat fparams'+  let inputs = do+        (p, i) <- zip fparams [0 ..]+        pure (paramName p, DistInput (ResTag i) (paramType p))+  let rettype' =+        addRetAls (map paramDeclType fparams'') $+          liftRetType w (map fst rettype)+  let (inputs', dstms) =+        distributeBody DistributeIrregular funHasParallelism const_scope [Var (paramName wp)] inputs body+      env = DistEnv $ M.fromList $ zip (map ResTag [0 ..]) reps+  -- Lift the body of the function and get the results, inserting copies as+  -- necessary to ensure the results are fresh and unique (see 'freshenResult').+  (body', needs) <-+    runFlattenM (castScope const_scope <> scopeOfFParams fparams'') $+      buildBody_ . freshenResult fparams'' $+        liftBody funHasParallelism funSizeParams defaultSegLevel w inputs' env dstms $+          bodyResult body+  let name = liftFunName $ funDefName fd+  pure+    ( fd+        { funDefName = name,+          funDefBody = body',+          funDefParams = fparams'',+          funDefRetType = rettype'+        },+      needs+    )++-- Here we assume that every type size is invariant and therefore every input+-- array is regular. As a result, parameters that correspond to type sizes are+-- not lifted and are also not part of 'DistInput'.+-- A uniformly lifted function can still return irregular arrays. This happens+-- when it returns an array whose dimension size was created in the function+-- body. In other words, the array has an existential size.+liftUniformFunDef ::+  FunHasParallelism ->+  FunSizeParams ->+  Scope SOACS ->+  FunDef SOACS ->+  PassM (FunDef GPU, S.Set DemandFn)+liftUniformFunDef funHasParallelism funSizeParams const_scope fd = do+  let FunDef+        { funDefBody = body,+          funDefParams = fparams,+          funDefRetType = rettype+        } = fd+  wp <- newParam "w" $ Prim int64+  let w = Var $ paramName wp+  let size_positions = funSizeParams $ funDefName fd+      isSizeParam = (`S.member` size_positions) . fst+      (indexed_sizes, indexed_values) =+        L.partition isSizeParam $ zip [0 ..] fparams+      fparam_sizes = map snd indexed_sizes+      fparams_explicit = map snd indexed_values++  (fparams_explicit', value_reps) <- mapAndUnzipM (liftRegularParam w) fparams_explicit+  let fparams'' = wp : fparam_sizes <> fparams_explicit'+  let inputs = do+        (p, i) <- zip fparams_explicit [0 ..]+        pure (paramName p, DistInput (ResTag i) (paramType p))+  let (inputs', dstms) =+        distributeBody DistributeIrregular funHasParallelism (const_scope <> scopeOfFParams fparam_sizes) [Var (paramName wp)] inputs body+      env = DistEnv $ M.fromList $ zip (map ResTag [0 ..]) value_reps+      rettype' =+        addRetAls (map paramDeclType fparams'') $+          liftRegularRetType inputs' w (map fst rettype)+  -- Lift the body of the function and get the results, inserting copies as+  -- necessary to ensure the results are fresh and unique (see 'freshenResult').+  (body', needs) <-+    runFlattenM (castScope const_scope <> scopeOfFParams fparams'') $+      buildBody_ . freshenResult fparams'' $+        -- XXX: I think function lifting makes it more important to classify invariant+        -- results in bodies. Function bodies can produce values that are+        -- invariant to the map-nest, but at this point there is no opportunity to+        -- hoist them out of the nest.++        liftUniformFunBody funHasParallelism funSizeParams defaultSegLevel w inputs' env dstms (map fst rettype) $+          bodyResult body+  let name = liftUniformFunName $ funDefName fd+  pure+    ( fd+        { funDefName = name,+          funDefBody = body',+          funDefParams = fparams'',+          funDefRetType = rettype'+        },+      needs+    )++transformLambda :: FunHasParallelism -> FunSizeParams -> Lambda SOACS -> FlattenM (Lambda GPU)+transformLambda funHasParallelism funSizeParams (Lambda params ret body) = do+  body' <- localScope (scopeOfLParams params) $ transformBody funHasParallelism funSizeParams body+  pure $ Lambda params ret body'++transformStm :: FunHasParallelism -> FunSizeParams -> Stm SOACS -> FlattenM ()+transformStm funHasParallelism funSizeParams (Let pat aux (Op soac))+  | "sequential_outer" `inAttrs` stmAuxAttrs aux = do+      scope <- askScope+      stms <- runBuilderT_ (FOT.transformSOAC pat soac) (castScope scope)+      transformStms funHasParallelism funSizeParams $ fmap (certify (stmAuxCerts aux)) stms+transformStm _ _ stm+  | "sequential" `inAttrs` stmAuxAttrs (stmAux stm) = addStm $ soacsStmToGPU stm+transformStm _ _ (Let pat aux (Op (Hist w arrs ops bucket_fun))) =+  certifying (stmAuxCerts aux) $ do+    res <-+      genUniformSegHist+        defaultSegLevel+        "topLevelSegHist"+        [w]+        ops+        (soacsLambdaToGPU bucket_fun)+        arrs+        (const $ pure ())+    forM_ (zip (patNames pat) res) $ \(v, v') ->+      letBindNames [v] $ BasicOp $ SubExp $ Var v'+transformStm funHasParallelism funSizeParams (Let pat aux (Op (Screma w arrs form)))+  | shouldDissectForm form =+      error "transformStm: complex Screma survived preprocessing"+  | otherwise =+      transformTopLevelScrema funHasParallelism funSizeParams pat aux w arrs form+transformStm funHasParallelism funSizeParams (Let pat aux (Op (FlatMap w arrs lam))) =+  certifying (stmAuxCerts aux) $ flattenFlatMap ops pat w arrs lam+  where+    irreg = irregularityFor DistributeIrregular aux+    ops = flattenOpsFor funHasParallelism funSizeParams irreg defaultSegLevel+transformStm funHasParallelism funSizeParams (Let pat aux (Loop params form body)) =+  localScope (scopeOfLoopForm form <> scopeOfFParams (map fst params)) $+    addStm . Let pat aux . Loop params form =<< transformBody funHasParallelism funSizeParams body+transformStm funHasParallelism funSizeParams (Let pat aux (Match ses cases def_body ret)) =+  addStm . Let pat aux+    =<< (Match ses <$> mapM onCase cases <*> transformBody funHasParallelism funSizeParams def_body <*> pure ret)+  where+    onCase = traverse (transformBody funHasParallelism funSizeParams)+transformStm funHasParallelism funSizeParams (Let pat aux (WithAcc inputs withacc_lam)) = do+  addStm . Let pat aux . WithAcc (map onInput inputs)+    =<< transformLambda funHasParallelism funSizeParams withacc_lam+  where+    onInput (shape, arrs, Nothing) =+      (shape, arrs, Nothing)+    onInput (shape, arrs, Just (lam, nes)) =+      (shape, arrs, Just (soacsLambdaToGPU lam, nes))+transformStm _ _ stm = addStm $ soacsStmToGPU stm++transformStms :: FunHasParallelism -> FunSizeParams -> Stms SOACS -> FlattenM ()+transformStms funHasParallelism funSizeParams stms =+  localScope (castScope $ scopeOf stms) $+    fold <$> traverse (transformStm funHasParallelism funSizeParams) stms++transformBody :: FunHasParallelism -> FunSizeParams -> Body SOACS -> FlattenM (Body GPU)+transformBody funHasParallelism funSizeParams (Body () stms res) = buildBody_ $ do+  transformStms funHasParallelism funSizeParams stms+  pure res++transformFunDef ::+  FunHasParallelism ->+  FunSizeParams ->+  Scope SOACS ->+  FunDef SOACS ->+  PassM (FunDef GPU, S.Set DemandFn)+transformFunDef funHasParallelism funSizeParams consts_scope fd = do+  let FunDef+        { funDefBody = body,+          funDefParams = fparams,+          funDefRetType = rettype+        } = fd+  (body', needs) <-+    runFlattenM (scopeOfFParams fparams <> castScope consts_scope) $+      transformBody funHasParallelism funSizeParams body+  pure+    ( fd+        { funDefBody = body',+          funDefRetType = rettype,+          funDefParams = fparams+        },+      needs+    )++liftUntilFixedPoint ::+  Prog SOACS ->+  FunHasParallelism ->+  FunSizeParams ->+  Scope SOACS ->+  S.Set DemandFn ->+  S.Set DemandFn ->+  PassM [FunDef GPU]+liftUntilFixedPoint prog funHasParallelism funSizeParams consts_scope made needed = do+  let made' = made <> needed+  (lifted_funs, new_needed) <-+    fmap (second ((`S.difference` made') . mconcat)) $+      mapAndUnzipM mkDemanded $+        S.toList needed+  if new_needed == mempty+    then pure lifted_funs+    else+      (lifted_funs ++)+        <$> liftUntilFixedPoint prog funHasParallelism funSizeParams consts_scope made' new_needed+  where+    mkDemanded (DemandLifted fname mode) =+      case find ((== fname) . funDefName) $ progFuns prog of+        Just fundef ->+          case mode of+            UniformLift -> liftUniformFunDef funHasParallelism funSizeParams consts_scope fundef+            NonUniformLift -> liftFunDef funHasParallelism funSizeParams consts_scope fundef+        Nothing -> error $ "mkDemanded: " <> show fname+    mkDemanded (DemandBuiltin b) = pure (builtinFunDef b, mempty)++transformProg :: Prog SOACS -> PassM (Prog GPU)+transformProg prog = do+  progAfterPreProcessing <- preprocessProg prog+  let consts = progConsts progAfterPreProcessing+      consts_scope = scopeOf consts+      funs = progFuns progAfterPreProcessing+      funParallelism = analyseFunParallelism funs+      size_param_map = analyseFunSizeParams funs+      funHasParallelism fname =+        M.findWithDefault (not $ isBuiltInFunction fname) fname funParallelism+      funSizeParams fname =+        M.findWithDefault mempty fname size_param_map+  (consts', consts_needs) <-+    runFlattenM mempty $ collectStms_ $ transformStms funHasParallelism funSizeParams consts+  (funs', funs_needs) <-+    second mconcat+      <$> mapAndUnzipM (transformFunDef funHasParallelism funSizeParams consts_scope) funs++  -- Now do fixpoint iteration until all needed functions have been provided.+  lifted_funs <-+    liftUntilFixedPoint+      prog+      funHasParallelism+      funSizeParams+      consts_scope+      mempty+      (consts_needs <> funs_needs)++  pure $+    prog+      { progConsts = consts',+        progFuns = lifted_funs <> funs'+      }++-- | Transform a SOACS program to a GPU program, using flattening.+flattenSOACs :: Pass SOACS GPU+flattenSOACs =+  Pass+    { passName = "flatten",+      passDescription = "Perform full flattening",+      passFunction = transformProg+    }+{-# NOINLINE flattenSOACs #-}
+ src/Futhark/Pass/Flatten/BasicOp.hs view
@@ -0,0 +1,1194 @@+module Futhark.Pass.Flatten.BasicOp (flattenBasicOp) where++import Control.Monad+import Data.Foldable+import Data.List qualified as L+import Data.List.NonEmpty qualified as NE+import Data.Tuple.Solo+import Futhark.IR.GPU+import Futhark.Pass.Flatten.Distribute+import Futhark.Pass.Flatten.General+import Futhark.Tools+import Futhark.Transform.Rename+import Futhark.Util.IntegralExp+import Prelude hiding (div, quot, rem)++-- Do 'map2 (++) A B' where 'A' and 'B' are irregular arrays and have the same+-- number of subarrays+concatIrreg ::+  SegLevel ->+  Segments ->+  DistEnv ->+  VName ->+  [IrregularRep] ->+  FlattenM IrregularRep+concatIrreg lvl _segments _env ns reparr = do+  -- Concatenation does not change the number of segments - it simply+  -- makes each of them larger.++  num_segments <- arraySize 0 <$> lookupType ns++  -- Constructs the full list size / shape that should hold the final results.+  ns_full <- letExp (baseName ns <> "_full") <=< segMap lvl (MkSolo num_segments) $+    \(MkSolo i) -> do+      old_segments <-+        forM reparr $ \rep ->+          letSubExp "old_segment" =<< eIndex (irregularS rep) [eSubExp i]+      new_segment <-+        letSubExp "new_segment"+          =<< toExp (sum $ map pe64 old_segments)+      pure $ subExpsRes [new_segment]++  (ns_full_F, ns_full_O, _ns_II1) <- doRepIota lvl ns_full++  repIota <- mapM (doRepIota lvl . irregularS) reparr+  segIota <- mapM (doSegIota lvl . irregularS) reparr++  let (_, _, rep_II1) = unzip3 repIota+  let (_, _, rep_II2) = unzip3 segIota++  n_arr <- mapM (fmap (arraySize 0) . lookupType) rep_II1++  -- Calculate offsets for the scatter operations+  let shapes = map irregularS reparr+  scatter_offsets <-+    letTupExp "irregular_scatter_offsets" <=< segMap lvl (MkSolo num_segments) $+      \(MkSolo i) -> do+        segment_sizes <-+          forM shapes $ \shape ->+            letSubExp "segment_size" =<< eIndex shape [eSubExp i]+        let scanned = scanl (+) 0 $ map pe64 segment_sizes+        sumprefix <- mapM (letSubExp "segment_prefix" <=< toExp) (init scanned)+        pure $ subExpsRes sumprefix++  scatter_offsets_T <-+    letTupExp "irregular_scatter_offsets_T" <=< segMap lvl (MkSolo num_segments) $+      \(MkSolo i) -> do+        columns <-+          forM scatter_offsets $ \offsets ->+            letSubExp "segment_offset" =<< eIndex offsets [eSubExp i]+        pure $ subExpsRes columns++  m <- arraySize 0 <$> lookupType ns_full_F+  data_t <- lookupType (irregularD (head reparr))+  let pt = elemType data_t+  let result_type = Array pt (Shape [m]) NoUniqueness+  elems_blank <- letExp "blank_res" =<< eBlank result_type++  -- Scatter data into result array+  elems <-+    foldlM+      ( \elems (reparr1, scatter_offset, n, ii1, ii2) -> do+          letExp "irregular_scatter_elems" <=< genScatter lvl elems n $ \gid -> do+            -- Which segment we are in.+            segment_i <-+              letSubExp "segment_i" =<< eIndex ii1 [eSubExp gid]++            -- Get segment offset in final array+            segment_o <-+              letSubExp "segment_o" =<< eIndex ns_full_O [eSubExp segment_i]++            -- Get local segment offset+            segment_local_o <-+              letSubExp "segment_local_o"+                =<< eIndex scatter_offset [eSubExp segment_i]++            o' <- letSubExp "o" =<< eIndex ii2 [eSubExp gid]+            src_segment_o <-+              letSubExp "src_segment_o" =<< eIndex (irregularO reparr1) [eSubExp segment_i]+            src_i <-+              letSubExp "src_i" <=< toExp $ pe64 src_segment_o + pe64 o'+            v' <-+              letSubExp "v" =<< eIndex (irregularD reparr1) [eSubExp src_i]++            -- Index to write `v'` at+            i <-+              letExp "i" =<< toExp (pe64 o' + pe64 segment_local_o + pe64 segment_o)++            pure (i, v')+      )+      elems_blank+      $ L.zip5 reparr scatter_offsets_T n_arr rep_II1 rep_II2++  pure $+    IrregularRep+      { irregularS = ns_full,+        irregularF = ns_full_F,+        irregularO = ns_full_O,+        irregularD = elems,+        irregularK = Dense+      }++-- We also can do reearange -> concat -> rearrange but this should be more efficient+concatIrregAlongDim ::+  SegLevel ->+  Segments ->+  DistEnv ->+  VName ->+  [IrregularRep] ->+  [Type] ->+  DistInputs ->+  Int ->+  FlattenM IrregularRep+concatIrregAlongDim lvl segments env ns rep_arr type_arr inps d = do+  num_segments <- arraySize 0 <$> lookupType ns++  ns_full <- letExp (baseName ns <> "_full") <=< segMap lvl (MkSolo num_segments) $+    \(MkSolo i) -> do+      old_segments <-+        forM rep_arr $ \rep ->+          letSubExp "old_segment" =<< eIndex (irregularS rep) [eSubExp i]+      new_segment <-+        letSubExp "new_segment"+          =<< toExp (sum $ map pe64 old_segments)+      pure $ subExpsRes [new_segment]++  (ns_full_F, ns_full_O, _ns_II1) <- doRepIota lvl ns_full++  repIota <- mapM (doRepIota lvl . irregularS) rep_arr+  segIota <- mapM (doSegIota lvl . irregularS) rep_arr++  let (_, _, rep_II1) = unzip3 repIota+  let (_, _, rep_II2) = unzip3 segIota++  n_arr <- mapM (fmap (arraySize 0) . lookupType) rep_II1++  scatter_info <-+    letTupExp "irregular_scatter_offsets" <=< segMap lvl (MkSolo num_segments) $+      \(MkSolo i) -> do+        seg_is <- segmentCoordsFromFlat segments i++        block_sizes <-+          forM type_arr $ \t -> do+            v_dims <- readTypeDims segments env seg_is inps t+            letSubExp "block_size" =<< toExp (product $ map pe64 $ drop d v_dims)++        let scanned = scanl (+) 0 $ map pe64 block_sizes+        sum_prefix <- mapM (letSubExp "segment_prefix" <=< toExp) (init scanned)+        total_block <- letSubExp "total_block" =<< toExp (last scanned)++        pure $ subExpsRes (block_sizes <> sum_prefix <> [total_block])++  let k = length type_arr+      (scatter_blocks, rest) = splitAt k scatter_info+      (scatter_offsets, [total_block_size]) = splitAt k rest++  m <- arraySize 0 <$> lookupType ns_full_F+  data_t <- lookupType (irregularD (head rep_arr))+  let pt = elemType data_t+  let result_type = Array pt (Shape [m]) NoUniqueness+  elems_blank <- letExp "blank_res" =<< eBlank result_type++  -- Scatter data into result array+  elems <-+    foldlM+      ( \elems (reparr1, scatter_block, scatter_offset, n, ii1, ii2) -> do+          letExp "irregular_scatter_elems" <=< genScatter lvl elems n $ \gid -> do+            -- Which segment we are in.+            segment_i <-+              letSubExp "segment_i" =<< eIndex ii1 [eSubExp gid]++            -- Get segment offset in final array+            segment_o <-+              letSubExp "segment_o" =<< eIndex ns_full_O [eSubExp segment_i]++            -- Get local segment offset+            segment_local_o <-+              letSubExp "segment_local_o"+                =<< eIndex scatter_offset [eSubExp segment_i]++            o' <- letSubExp "o" =<< eIndex ii2 [eSubExp gid]+            src_segment_o <-+              letSubExp "src_segment_o" =<< eIndex (irregularO reparr1) [eSubExp segment_i]+            src_i <-+              letSubExp "src_i" <=< toExp $ pe64 src_segment_o + pe64 o'+            v' <-+              letSubExp "v" =<< eIndex (irregularD reparr1) [eSubExp src_i]++            scatter_block_size <-+              letSubExp "scatter_block_size" =<< eIndex scatter_block [eSubExp segment_i]++            scatter_total_block_size <-+              letSubExp "scatter_total_block_size" =<< eIndex total_block_size [eSubExp segment_i]++            outer_i <-+              letSubExp "outer_i" =<< toExp (pe64 o' `div` pe64 scatter_block_size)++            i <-+              letExp "i"+                =<< toExp+                  ( pe64 o'+                      + pe64 outer_i * (pe64 scatter_total_block_size - pe64 scatter_block_size)+                      + pe64 segment_local_o+                      + pe64 segment_o+                  )+            pure (i, v')+      )+      elems_blank+      $ L.zip6 rep_arr scatter_blocks scatter_offsets n_arr rep_II1 rep_II2++  pure $+    IrregularRep+      { irregularS = ns_full,+        irregularF = ns_full_F,+        irregularO = ns_full_O,+        irregularD = elems,+        irregularK = Dense+      }++-- Do 'map2 replicate ns A', where 'A' is an irregular array (and so+-- is the result, obviously).+replicateIrreg ::+  SegLevel ->+  Segments ->+  DistEnv ->+  VName ->+  Name ->+  IrregularRep ->+  FlattenM IrregularRep+replicateIrreg lvl _segments _env ns desc rep = do+  -- Replication does not change the number of segments - it simply+  -- makes each of them larger.++  num_segments <- arraySize 0 <$> lookupType ns++  -- ns multipled with existing segment sizes.+  ns_full <- letExp (baseName ns <> "_full") <=< segMap lvl (MkSolo num_segments) $+    \(MkSolo i) -> do+      n <-+        letSubExp "n" =<< eIndex ns [eSubExp i]+      old_segment <-+        letSubExp "old_segment" =<< eIndex (irregularS rep) [eSubExp i]+      full_segment <-+        letSubExp "new_segment" =<< toExp (pe64 n * pe64 old_segment)+      pure $ subExpsRes [full_segment]++  (ns_full_F, ns_full_O, ns_full_D) <- doRepIota lvl ns_full+  (_, _, flat_to_segs) <- doSegIota lvl ns_full++  w <- arraySize 0 <$> lookupType ns_full_D++  elems <- letExp (desc <> "_rep_D") <=< segMap lvl (MkSolo w) $ \(MkSolo i) -> do+    -- Which segment we are in.+    segment_i <-+      letSubExp "segment_i" =<< eIndex ns_full_D [eSubExp i]+    -- Size of original segment.+    old_segment <-+      letSubExp "old_segment" =<< eIndex (irregularS rep) [eSubExp segment_i]+    -- Index of value inside *new* segment.+    j_new <-+      letSubExp "j_new" =<< eIndex flat_to_segs [eSubExp i]+    -- Index of value inside *old* segment.+    j_old <-+      letSubExp "j_old" =<< toExp (pe64 j_new `rem` pe64 old_segment)+    -- Offset of values in original segment.+    offset <-+      letSubExp "offset" =<< eIndex (irregularO rep) [eSubExp segment_i]+    v <-+      letSubExp "v"+        =<< eIndex (irregularD rep) [toExp $ pe64 offset + pe64 j_old]+    pure $ subExpsRes [v]++  pure $+    IrregularRep+      { irregularS = ns_full,+        irregularF = ns_full_F,+        irregularO = ns_full_O,+        irregularD = elems,+        irregularK = Dense+      }++rearrangeFlat :: (IntegralExp num) => [Int] -> [num] -> num -> num+rearrangeFlat perm dims i =+  flattenIndex dims $+    rearrangeShape (rearrangeInverse perm) $+      unflattenIndex (rearrangeShape perm dims) i++segmentCoordsFromFlat :: Segments -> SubExp -> FlattenM [SubExp]+segmentCoordsFromFlat segments seg_i =+  mapM (letSubExp "seg_coord" <=< toExp) $+    unflattenIndex (map pe64 $ shapeDims $ segmentsShape segments) (pe64 seg_i)++-- TODO: We do not need to actually make this Dense+rearrangeIrreg ::+  SegLevel ->+  Segments ->+  DistEnv ->+  DistInputs ->+  TypeBase Shape u ->+  [Int] ->+  IrregularRep ->+  FlattenM IrregularRep+rearrangeIrreg lvl segments env inps v_t perm ir = do+  (IrregularRep shape _ offsets elems _) <- flattenIrregularRep lvl ir+  (new_F, new_O, ii1_vss) <- doRepIota lvl shape+  (_, _, ii2_vss) <- doSegIota lvl shape+  m <- arraySize 0 <$> lookupType ii1_vss+  elems' <- letExp "elems_rearrange" <=< renameExp <=< segMap lvl (MkSolo m) $+    \(MkSolo i) -> do+      seg_i <- letSubExp "seg_i" =<< eIndex ii1_vss [eSubExp i]+      offset <- letSubExp "offset" =<< eIndex offsets [eSubExp seg_i]+      in_seg_i <- letSubExp "in_seg_i" =<< eIndex ii2_vss [eSubExp i]+      seg_is <- segmentCoordsFromFlat segments seg_i+      v_dims <- readTypeDims segments env seg_is inps v_t+      let v_dims' = map pe64 v_dims+          in_seg_is_tr = rearrangeFlat perm v_dims' $ pe64 in_seg_i+      v' <-+        letSubExp "v"+          =<< eIndex elems [toExp $ pe64 offset + in_seg_is_tr]+      pure [subExpRes v']+  pure $+    IrregularRep+      { irregularS = shape,+        irregularF = new_F,+        irregularO = new_O,+        irregularD = elems',+        irregularK = Dense+      }++-- | Input common to 'flattenBasicOp': the segment level, the enclosing+-- segments, the distribution environment and inputs, the result being produced,+-- and the statement auxiliary information.+data TrCtx = TrCtx SegLevel Segments DistEnv DistInputs DistResult (StmAux ())++transformArrayLit ::+  TrCtx ->+  [SubExp] ->+  Type ->+  FlattenM DistEnv+-- Potentially no need for this case.+transformArrayLit (TrCtx lvl segments env inps res _aux) [] row_type+  | not $ any (isVariant inps) (arrayDims row_type) = do+      let resultType =+            Array+              (elemType row_type)+              (segmentsShape segments <> Shape [intConst Int64 0] <> arrayShape row_type)+              NoUniqueness+      v <- letExp "arraylit_empty_reg" =<< eBlank resultType+      pure $ insertRegulars [distResTag res] [v] env+  | otherwise = do+      ns <- dataArr lvl segments env inps $ intConst Int64 0+      (flags, offsets, _elems) <- doRepIota lvl ns+      let resultType = Array (elemType row_type) (Shape [intConst Int64 0]) NoUniqueness+      elems <- letExp "arraylit_empty_elems" =<< eBlank resultType+      insertIrregularM ns flags offsets (distResTag res) elems Dense env+transformArrayLit (TrCtx lvl segments env inps res _aux) vs row_type+  | not $ any (isVariant inps) (arrayDims row_type) = do+      res_v <-+        if any (isVariant inps) vs+          then do+            let seg_shape = segmentsShape segments+                one = intConst Int64 1+                arr_outer_dim = intConst Int64 $ toInteger $ length vs+                expected = seg_shape <> arrayShape row_type+                stacked = seg_shape <> Shape [one] <> arrayShape row_type+                d = segmentsRank segments++            vs_reg <- mapM (liftSubExpRegular lvl segments inps env expected) vs++            vs_reg_1 <-+              forM vs_reg $ \v -> do+                v_t <- lookupType v+                letExp (baseName v <> "_stack") $+                  BasicOp $+                    Reshape v $+                      reshapeAll (arrayShape v_t) stacked++            case vs_reg_1 of+              [] -> error "flattenBasicOp: empty ArrayLit cannot have variant elements"+              [v] ->+                pure v+              v : vs' ->+                letExp "arraylit_reg" $ BasicOp $ Concat d (v NE.:| vs') arr_outer_dim+          else do+            base_v <- letExp "arraylit_base" $ BasicOp $ ArrayLit vs row_type+            letExp "arraylit_reg" $+              BasicOp $+                Replicate (segmentsShape segments) (Var base_v)+      pure $ insertRegulars [distResTag res] [res_v] env+  | otherwise = do+      let arr_outer_dim = intConst Int64 $ fromIntegral $ length vs+      vs_reparr <- mapM (dataArr lvl segments env inps) vs+      dim_arrs <- mapM (dataArr lvl segments env inps) (arrayDims row_type)+      num_segments <- letSubExp "num_segments" =<< toExp (segmentCount segments)+      ~[row_size, full_size] <- letTupExp "arraylit_row_size" <=< segMap lvl (MkSolo num_segments) $ \(MkSolo i) -> do+        vals <- mapM (\dim_arr -> letSubExp "dim_i" =<< eIndex dim_arr [eSubExp i]) dim_arrs+        n <- letSubExp "n" <=< toExp $ product $ map pe64 vals+        fs <- letSubExp "fs" <=< toExp $ pe64 n * pe64 arr_outer_dim+        pure $ subExpsRes [n, fs]++      (_, _, row_II1) <- doRepIota lvl row_size+      (_, _, row_II2) <- doSegIota lvl row_size++      row_flat_size <- arraySize 0 <$> lookupType row_II1++      (full_flags, full_offset, full_II1) <- doRepIota lvl full_size++      m <- arraySize 0 <$> lookupType full_II1+      let pt = elemType row_type+      let resultType = Array pt (Shape [m]) NoUniqueness+      elems_blank <- letExp "blank_res" =<< eBlank resultType++      elems <-+        foldlM+          ( \elems (var_num, arr) -> do+              letExp "irregular_scatter_elems" <=< genScatter lvl elems row_flat_size $ \gid -> do+                -- Which segment we are in.+                segment_i <-+                  letSubExp "segment_i" =<< eIndex row_II1 [eSubExp gid]++                row_size_i <-+                  letSubExp "row_size_i" =<< eIndex row_size [eSubExp segment_i]++                segment_global_o <-+                  letSubExp "segment_global_o"+                    =<< eIndex full_offset [eSubExp segment_i]++                v' <-+                  letSubExp "v" =<< eIndex arr [eSubExp gid]++                o' <- letSubExp "o" =<< eIndex row_II2 [eSubExp gid]++                i <-+                  letExp "i"+                    =<< toExp+                      ( pe64 o'+                          + pe64 segment_global_o+                          + pe64 row_size_i * pe64 (intConst Int64 var_num)+                      )++                pure (i, v')+          )+          elems_blank+          $ zip [0 ..] vs_reparr++      insertIrregularM full_size full_flags full_offset (distResTag res) elems Dense env++transformArrayVal ::+  TrCtx ->+  [PrimValue] ->+  PrimType ->+  FlattenM DistEnv+transformArrayVal (TrCtx _lvl segments env _inps res _aux) vs row_type = do+  base_v <- letExp "arraylit_base" $ BasicOp $ ArrayVal vs row_type+  res_v <- letExp "arraylit_reg" $ BasicOp $ Replicate (segmentsShape segments) (Var base_v)+  pure $ insertRegulars [distResTag res] [res_v] env++transformReshape ::+  TrCtx ->+  VName ->+  NewShape SubExp ->+  FlattenM DistEnv+transformReshape (TrCtx lvl segments env inps res aux) arr reshape+  | isRegularDistResult res,+    not (any (isVariant inps) reshape) = do+      let outer = segmentsShape segments+          inner_target = newShape reshape+          reshape' = reshapeCoerce outer <> newshapeInner outer reshape++      arr_t <- lookupInputType inps arr+      let arr_shape = arrayShape arr_t+      let unform_arr = not (any (isVariant inps) arr_shape)+      if unform_arr+        then do+          arr' <-+            liftSubExpRegular+              lvl+              segments+              inps+              env+              (outer <> arr_shape)+              (Var arr)+          v <- certifying (distCerts inps aux env) . letExp "reshape_reg" . BasicOp $ Reshape arr' reshape'+          pure $ insertRegulars [distResTag res] [v] env+        else do+          arr' <-+            liftSubExpRegular+              lvl+              segments+              inps+              env+              (outer <> inner_target)+              (Var arr)+          pure $ insertRegulars [distResTag res] [arr'] env+  | otherwise = do+      irreg_v <- getIrregRep lvl segments env inps arr+      insertRepM (distResTag res) (Irregular irreg_v) env++transformIndex ::+  TrCtx ->+  VName ->+  Slice SubExp ->+  FlattenM DistEnv+transformIndex (TrCtx lvl segments env inps res aux) arr slice+  | -- Must be a regular result...+    isRegularDistResult res,+    -- And a regular input...+    isRegularInputArr env inps arr,+    -- And uniform slices.+    not (any (isVariant inps) slice) = do+      arr_t <- lookupInputType inps arr+      arr' <-+        liftSubExpRegular+          lvl+          segments+          inps+          env+          (segmentsShape segments <> arrayShape arr_t)+          (Var arr)+      let segmentSlice = map sliceDim . shapeDims . segmentsShape+      v <-+        certifying (distCerts inps aux env) . letExp "index_reg" . BasicOp $+          Index arr' (Slice $ segmentSlice segments <> unSlice slice)+      pure $ insertRegulars [distResTag res] [v] env+  | isRegularDistResult res,+    not (any (isVariant inps) (sliceDims slice)) = do+      let space = shapeDims (segmentsShape segments) <> sliceDims slice+      v <-+        letExp "index_reg_gather"+          <=< renameExp+          <=< segMap lvl (NE.fromList space)+          $ \is -> do+            let (seg_is, in_is) = splitAt (segmentsRank segments) (toList is)+            readInputs segments env seg_is inps+            let slice' = fixSlice (fmap pe64 slice) (map pe64 in_is)+            auxing aux $+              fmap (subExpsRes . pure) . letSubExp "v"+                =<< eIndex arr (map toExp slice')+      pure $ insertRegulars [distResTag res] [v] env+  | otherwise = do+      -- Maximally nonuniform case.+      num_segments <- letSubExp "num_segments" =<< toExp (segmentCount segments)+      ns <- letExp "slice_sizes" <=< renameExp <=< segMap lvl (MkSolo num_segments) $ \(MkSolo segment) -> do+        segment_is <- segmentCoordsFromFlat segments segment+        slice_ns <- mapM (readInput segments env segment_is inps) $ sliceDims slice+        fmap varsRes . letTupExp "n" <=< toExp $ product $ map pe64 slice_ns+      (_n, offsets, m) <- exScanAndSum lvl ns+      (_, _, repiota_D) <- doRepIota lvl ns+      flags <- genFlags lvl m offsets+      elems <- letExp "index_irreg_elems" <=< renameExp <=< segMap lvl (NE.singleton m) $ \is -> do+        segment <- letSubExp "segment" =<< eIndex repiota_D (toList $ fmap eSubExp is)+        segment_start <- letSubExp "segment_start" =<< eIndex offsets [eSubExp segment]+        segment_is <- segmentCoordsFromFlat segments segment+        readInputs segments env segment_is inps+        let slice' =+              fixSlice (fmap pe64 slice) $+                unflattenIndex (map pe64 (sliceDims slice)) $+                  subtract (pe64 segment_start) . pe64 $+                    NE.head is+        auxing aux $+          fmap (subExpsRes . pure) . letSubExp "v"+            =<< eIndex arr (map toExp slice')+      insertIrregularM ns flags offsets (distResTag res) elems Dense env++transformFlatIndex ::+  TrCtx ->+  VName ->+  FlatSlice SubExp ->+  FlattenM DistEnv+transformFlatIndex (TrCtx lvl segments env inps res aux) arr flat_slice+  | isRegularDistResult res,+    not (any (isVariant inps) flat_slice) = do+      arr_t <- lookupInputType inps arr+      -- arr should be 1D+      let [n] = arrayDims arr_t+      num_segments <- letSubExp "num_segments" =<< toExp (segmentCount segments)+      arr_flat_size <- letSubExp "arr_flat_size" =<< toExp (pe64 num_segments * pe64 n)+      let arr_lift_shape = segmentsShape segments <> arrayShape arr_t+          arr_flat_shape = Shape [arr_flat_size]+      arr' <-+        liftSubExpRegular+          lvl+          segments+          inps+          env+          arr_lift_shape+          (Var arr)+      arr'_flat <-+        letExp (baseName arr <> "_reshaped") $ BasicOp $ Reshape arr' $ reshapeAll arr_lift_shape arr_flat_shape+      let FlatSlice off dims = flat_slice+          flat_slice' = FlatSlice off (FlatDimIndex num_segments n : dims)+      out_flat_updated <-+        certifying (distCerts inps aux env) . letExp "flat_index_reg" . BasicOp $+          FlatIndex arr'_flat flat_slice'+      out_updated <-+        letExp "flat_index_reg_reshaped" $+          BasicOp $+            Reshape out_flat_updated $+              reshapeAll arr_flat_shape arr_lift_shape+      pure $ insertRegulars [distResTag res] [out_updated] env+  | isRegularDistResult res,+    not (any (isVariant inps) (flatSliceDims flat_slice)) = do+      let space = shapeDims (segmentsShape segments) <> flatSliceDims flat_slice+      v <-+        letExp "flat_index_reg_gather"+          <=< renameExp+          <=< segMap lvl (NE.fromList space)+          $ \is -> do+            let (seg_is, in_is) = splitAt (segmentsRank segments) (toList is)+            readInputs segments env seg_is inps+            let flat_slice'@(FlatSlice flat_offset _) = fmap pe64 flat_slice+                flat_i =+                  flat_offset+                    + sum (zipWith (*) (map pe64 in_is) (flatSliceStrides flat_slice'))+            auxing aux $+              fmap (subExpsRes . pure) . letSubExp "v"+                =<< eIndex arr [toExp flat_i]+      pure $ insertRegulars [distResTag res] [v] env+  | otherwise = do+      -- Maximally nonuniform case.+      num_segments <- letSubExp "num_segments" =<< toExp (segmentCount segments)+      ns <- letExp "slice_sizes" <=< renameExp <=< segMap lvl (MkSolo num_segments) $ \(MkSolo segment) -> do+        segment_is <- segmentCoordsFromFlat segments segment+        slice_ns <- mapM (readInput segments env segment_is inps) $ flatSliceDims flat_slice+        fmap varsRes . letTupExp "n" <=< toExp $ product $ map pe64 slice_ns+      (_n, offsets, m) <- exScanAndSum lvl ns+      (_, _, repiota_D) <- doRepIota lvl ns+      flags <- genFlags lvl m offsets+      elems <- letExp "flat_index_irreg_elems" <=< renameExp <=< segMap lvl (NE.singleton m) $ \is -> do+        segment <- letSubExp "segment" =<< eIndex repiota_D (toList $ fmap eSubExp is)+        segment_start <- letSubExp "segment_start" =<< eIndex offsets [eSubExp segment]+        segment_is <- segmentCoordsFromFlat segments segment+        readInputs segments env segment_is inps+        let flat_slice'@(FlatSlice flat_offset _) = fmap pe64 flat_slice+            local_flat = pe64 (NE.head is) - pe64 segment_start+            local_is = unflattenIndex (flatSliceDims flat_slice') local_flat+            flat_i = flat_offset + sum (zipWith (*) local_is (flatSliceStrides flat_slice'))+        auxing aux $+          fmap (subExpsRes . pure) . letSubExp "v"+            =<< eIndex arr [toExp flat_i]+      insertIrregularM ns flags offsets (distResTag res) elems Dense env++transformIota ::+  TrCtx ->+  SubExp ->+  SubExp ->+  SubExp ->+  IntType ->+  FlattenM DistEnv+transformIota (TrCtx _lvl segments env inps res _aux) n x s it+  | isRegularDistResult res,+    not (isVariant inps n),+    not (isVariant inps x),+    not (isVariant inps s) = do+      iota_row <- letExp "iota_reg_row" $ BasicOp $ Iota n x s it+      v <-+        letExp "iota_reg" $+          BasicOp $+            Replicate (segmentsShape segments) (Var iota_row)+      pure $ insertRegulars [distResTag res] [v] env+transformIota (TrCtx lvl segments env inps res aux) n (Constant x) (Constant s) Int64+  | zeroIsh x,+    oneIsh s = do+      ns <- dataArr lvl segments env inps n+      (flags, offsets, elems) <- certifying (distCerts inps aux env) $ doSegIota lvl ns+      insertIrregularM ns flags offsets (distResTag res) elems Dense env+transformIota (TrCtx lvl segments env inps res aux) n x s it = do+  ns <- dataArr lvl segments env inps n+  xs <- dataArr lvl segments env inps x+  ss <- dataArr lvl segments env inps s+  (res_F, res_O, res_D) <- certifying (distCerts inps aux env) $ doSegIota lvl ns+  (_, _, repiota_D) <- doRepIota lvl ns+  m <- arraySize 0 <$> lookupType res_D+  res_D' <- letExp "iota_D_fixed" <=< segMap lvl (MkSolo m) $ \(MkSolo i) -> do+    segment <- letSubExp "segment" =<< eIndex repiota_D [eSubExp i]+    v' <- letSubExp "v" =<< eIndex res_D [eSubExp i]+    x' <- letSubExp "x" =<< eIndex xs [eSubExp segment]+    s' <- letSubExp "s" =<< eIndex ss [eSubExp segment]+    fmap (subExpsRes . pure) . letSubExp "v" <=< toExp $+      primExpFromSubExp (IntType it) x'+        ~+~ sExt it (untyped (pe64 v'))+        ~*~ primExpFromSubExp (IntType it) s'+  insertIrregularM ns res_F res_O (distResTag res) res_D' Dense env++transformConcat ::+  TrCtx ->+  Int ->+  NE.NonEmpty VName ->+  SubExp ->+  FlattenM DistEnv+transformConcat (TrCtx lvl segments env inps res _aux) d arr shp = do+  arr_ts <- mapM (lookupInputType inps) (NE.toList arr)+  let inputShapeUniform t =+        not $ any (isVariant inps) (arrayDims t)+  if isRegularDistResult res+    && not (isVariant inps shp)+    && all inputShapeUniform arr_ts+    then do+      --  Unifrom Concat+      arrs_lifted <-+        forM (zip (NE.toList arr) arr_ts) $ \(v, t) -> do+          let expectedShape = segmentsShape segments <> arrayShape t+          liftSubExpRegular lvl segments inps env expectedShape (Var v)+      v' <-+        letExp "concat_reg" $+          BasicOp $+            Concat+              (segmentsRank segments + d)+              (NE.fromList arrs_lifted)+              shp++      pure $ insertRegulars [distResTag res] [v'] env+    else do+      ns <- dataArr lvl segments env inps shp+      reparr <- mapM (getIrregRep lvl segments env inps) (NE.toList arr)+      rep' <- case d of+        0 -> concatIrreg lvl segments env ns reparr+        d' -> do+          concatIrregAlongDim lvl segments env ns reparr arr_ts inps d'+      insertRepM (distResTag res) (Irregular rep') env++transformReplicate ::+  TrCtx ->+  Shape ->+  SubExp ->+  FlattenM DistEnv+--  Uniform Replicate+transformReplicate (TrCtx lvl segments env inps res _aux) (Shape dims) se+  | isRegularDistResult res = do+      t <- subExpInputType inps se+      let expectedShape = segmentsShape segments <> arrayShape t+      lifted <- liftSubExpRegular lvl segments inps env expectedShape se+      v_rep <- replicateForDims segments (Shape dims) lifted+      pure $ insertRegulars [distResTag res] [v_rep] env+transformReplicate (TrCtx lvl segments env inps res _aux) (Shape [n]) (Var v) = do+  ns <- dataArr lvl segments env inps n+  rep <- getIrregRep lvl segments env inps v+  rep' <- replicateIrreg lvl segments env ns (baseName v) rep+  insertRepM (distResTag res) (Irregular rep') env+transformReplicate (TrCtx lvl segments env inps res aux) (Shape [n]) (Constant v) = do+  ns <- dataArr lvl segments env inps n+  (res_F, res_O, res_D) <-+    certifying (distCerts inps aux env) $ doSegIota lvl ns+  w <- arraySize 0 <$> lookupType res_D+  res_D' <- letExp "rep_const" $ BasicOp $ Replicate (Shape [w]) (Constant v)+  insertIrregularM ns res_F res_O (distResTag res) res_D' Dense env+transformReplicate (TrCtx lvl segments env inps res aux) (Shape dims) (Constant v) = do+  dim_arrs <- mapM (dataArr lvl segments env inps) dims+  seg_number <- arraySize 0 <$> lookupType (head dim_arrs)+  mul_dims <- letExp "mul_dims" <=< segMap lvl (MkSolo seg_number) $ \(MkSolo i) -> do+    vals <- mapM (\dim_arr -> letSubExp "dim_i" =<< eIndex dim_arr [eSubExp i]) dim_arrs+    n <- letSubExp "n" <=< toExp $ product $ map pe64 vals+    pure [subExpRes n]+  (res_F, res_O, res_D) <-+    certifying (distCerts inps aux env) $ doSegIota lvl mul_dims+  w <- arraySize 0 <$> lookupType res_D+  res_D' <- letExp "rep_const" $ BasicOp $ Replicate (Shape [w]) (Constant v)+  insertIrregularM mul_dims res_F res_O (distResTag res) res_D' Dense env+transformReplicate (TrCtx _lvl segments env inps res _aux) (Shape []) (Var v) =+  case lookup v inps of+    Just (DistInputFree v' _) -> do+      v'' <-+        letExp (baseName v' <> "_copy") . BasicOp $+          Replicate mempty (Var v')+      pure $ insertRegulars [distResTag res] [v''] env+    Just (DistInput rt _) ->+      case resVar rt env of+        Irregular r -> do+          let name = baseName (irregularD r) <> "_copy"+          elems_copy <-+            letExp name . BasicOp $+              Replicate mempty (Var $ irregularD r)+          let rep = Irregular $ r {irregularD = elems_copy}+          insertRepM (distResTag res) rep env+        Regular v' -> do+          v'' <-+            letExp (baseName v' <> "_copy") . BasicOp $+              Replicate mempty (Var v')+          pure $ insertRegulars [distResTag res] [v''] env+    Nothing -> do+      v' <-+        letExp (baseName v <> "_copy_free") . BasicOp $+          Replicate (segmentsShape segments) (Var v)+      pure $ insertRegulars [distResTag res] [v'] env+transformReplicate (TrCtx lvl segments env inps res _aux) (Shape dims) (Var v) = do+  dim_arrs <- mapM (dataArr lvl segments env inps) dims+  seg_number <- arraySize 0 <$> lookupType (head dim_arrs)+  mul_dims <- letExp "mul_dims" <=< segMap lvl (MkSolo seg_number) $ \(MkSolo i) -> do+    vals <- mapM (\dim_arr -> letSubExp "dim_i" =<< eIndex dim_arr [eSubExp i]) dim_arrs+    n <- letSubExp "n" <=< toExp $ product $ map pe64 vals+    pure [subExpRes n]+  rep <- getIrregRep lvl segments env inps v+  rep' <- replicateIrreg lvl segments env mul_dims (baseName v) rep+  insertRepM (distResTag res) (Irregular rep') env++transformManifest ::+  TrCtx ->+  VName ->+  [Int] ->+  FlattenM DistEnv+transformManifest (TrCtx lvl segments env inps res _aux) v perm+  | isRegularDistResult res = do+      t <- lookupInputType inps v+      v_lifted <-+        liftSubExpRegular+          lvl+          segments+          inps+          env+          (segmentsShape segments <> arrayShape t)+          (Var v)+      let segment_rank = segmentsRank segments+      v_manifest <- letExp (baseName v <> "_manifest") . BasicOp $ Manifest v_lifted ([0 .. segment_rank - 1] ++ map (+ segment_rank) perm)+      pure $ insertRegulars [distResTag res] [v_manifest] env+  | otherwise = do+      irreg <- getIrregRep lvl segments env inps v+      irreg_dense <- ensureDenseIrregular lvl (baseName v <> "_manifest") irreg+      elems_copy <-+        letExp (baseName (irregularD irreg_dense) <> "_manifest") . BasicOp $+          Replicate mempty (Var $ irregularD irreg_dense)+      insertRepM+        (distResTag res)+        (Irregular $ irreg_dense {irregularD = elems_copy})+        env++transformUpdate ::+  TrCtx ->+  Safety ->+  VName ->+  Slice SubExp ->+  SubExp ->+  FlattenM DistEnv+transformUpdate (TrCtx lvl segments env inps res aux) safety as slice se+  -- Uniform Update+  | Just as_t <- distInputType <$> lookup as inps,+    isRegularDistResult res,+    not (any (isVariant inps) slice) = do+      as' <-+        liftSubExpRegular+          lvl+          segments+          inps+          env+          (segmentsShape segments <> arrayShape as_t)+          (Var as)+      se' <-+        liftSubExpRegular+          lvl+          segments+          inps+          env+          (segmentsShape segments <> sliceShape slice)+          se+      let segmentSlice = map sliceDim . shapeDims . segmentsShape+      v <-+        certifying (distCerts inps aux env) . letExp "update_reg" . BasicOp $+          Update safety as' (Slice $ segmentSlice segments <> unSlice slice) (Var se')+      pure $ insertRegulars [distResTag res] [v] env+  | Just as_t <- distInputType <$> lookup as inps,+    isRegularDistResult res,+    not (any (isVariant inps) (sliceDims slice)) = do+      let as_lift_shape = segmentsShape segments <> arrayShape as_t+      as' <- liftSubExpRegular lvl segments inps env as_lift_shape (Var as)+      let update_dims = segments <> sliceDims slice+      updated <-+        certifying (distCerts inps aux env)+          . letExp "update_reg_scatter"+          <=< renameExp+          <=< genScatterND lvl as' update_dims+          $ \is -> do+            let (seg_is, in_is) = splitAt (segmentsRank segments) (toList is)+            readInputs segments env seg_is $ filter ((/= as) . fst) inps+            let slice' = fixSlice (fmap pe64 slice) (map pe64 in_is)+            -- Value to write+            v' <- case se of+              Constant c -> pure $ Constant c+              Var se_v -> letSubExp "v" =<< eIndex se_v (map toExp in_is)+            -- Index to write `v'` at+            in_is' <- mapM (letSubExp "i" <=< toExp) slice'+            let is' = seg_is <> in_is'+            pure (is', v')+      pure $ insertRegulars [distResTag res] [updated] env+  | Just as_t <- distInputType <$> lookup as inps = do+      num_segments <- letSubExp "num_segments" =<< toExp (segmentCount segments)+      ns <- letExp "slice_sizes"+        <=< renameExp+        <=< segMap lvl (MkSolo num_segments)+        $ \(MkSolo seg_i) -> do+          seg_is <- segmentCoordsFromFlat segments seg_i+          readInputs segments env seg_is $+            filter ((`elem` sliceDims slice) . Var . fst) inps+          slice_dims <- mapM (readInput segments env seg_is inps) $ sliceDims slice+          n <- letSubExp "n" <=< toExp $ product $ map pe64 slice_dims+          pure [subExpRes n]+      -- Irregular representation of `as`+      as_rep <- getIrregRep lvl segments env inps as+      IrregularRep shape flags offsets elems _ <-+        ensureDenseIrregular lvl (baseName as <> "_update") as_rep+      -- Inner indices (1 and 2) of `ns`+      (_, _, ii1_vss) <- doRepIota lvl ns+      (_, _, ii2_vss) <- certifying (distCerts inps aux env) $ doSegIota lvl ns+      -- Number of updates to perform+      m <- arraySize 0 <$> lookupType ii2_vss+      elems' <- letExp "elems_scatter" <=< renameExp <=< genScatter lvl elems m $ \gid -> do+        seg_i <- letSubExp "seg_i" =<< eIndex ii1_vss [eSubExp gid]+        in_seg_i <- letSubExp "in_seg_i" =<< eIndex ii2_vss [eSubExp gid]+        seg_is <- segmentCoordsFromFlat segments seg_i+        readInputs segments env seg_is $ filter ((/= as) . fst) inps+        as_dims <- readTypeDims segments env seg_is inps as_t+        slice_dims <- mapM (readInput segments env seg_is inps) $ sliceDims slice+        case se of+          Var v -> do+            let in_seg_is =+                  unflattenIndex (map pe64 slice_dims) (pe64 in_seg_i)+                slice' = fmap pe64 slice+                flat_i =+                  flattenIndex+                    (map pe64 as_dims)+                    (fixSlice slice' in_seg_is)+            -- Value to write+            v' <- letSubExp "v" =<< eIndex v (map toExp in_seg_is)+            o' <- letSubExp "o" =<< eIndex offsets [eSubExp seg_i]+            -- Index to write `v'` at+            i <- letExp "i" =<< toExp (pe64 o' + flat_i)+            pure (i, v')+          Constant c -> do+            let slice' = fmap pe64 slice+                flat_i = flattenIndex (map pe64 as_dims) (fixSlice slice' [])+            o' <- letSubExp "o" =<< eIndex offsets [eSubExp seg_i]+            i <- letExp "i" =<< toExp (pe64 o' + flat_i)+            pure (i, Constant c)+      insertIrregularM shape flags offsets (distResTag res) elems' Dense env+  | otherwise =+      error "Flattening update: destination is not input."++transformFlatUpdate ::+  TrCtx ->+  VName ->+  FlatSlice SubExp ->+  VName ->+  FlattenM DistEnv+transformFlatUpdate (TrCtx lvl segments env inps res aux) as flat_slice v+  -- Uniform Flat Update+  | Just as_t <- distInputType <$> lookup as inps,+    isRegularDistResult res,+    not (any (isVariant inps) flat_slice) = do+      -- as should be 1D+      let [n] = arrayDims as_t+      num_segments <- letSubExp "num_segments" =<< toExp (segmentCount segments)+      as_flat_size <- letSubExp "as_flat_size" =<< toExp (pe64 num_segments * pe64 n)++      let se_shape = Shape $ flatSliceDims flat_slice+          as_lift_shape = segmentsShape segments <> arrayShape as_t+          se_lift_shape = segmentsShape segments <> se_shape+          se_flat_shape = Shape [num_segments] <> se_shape+          as_flat_shape = Shape [as_flat_size]+      as' <-+        liftSubExpRegular+          lvl+          segments+          inps+          env+          as_lift_shape+          (Var as)+      v' <-+        liftSubExpRegular+          lvl+          segments+          inps+          env+          se_lift_shape+          (Var v)+      as_flat <-+        letExp (baseName as <> "_reshaped") $ BasicOp $ Reshape as' $ reshapeAll as_lift_shape as_flat_shape+      v_flat <-+        letExp (baseName v <> "_reshaped") $ BasicOp $ Reshape v' $ reshapeAll se_lift_shape se_flat_shape+      let FlatSlice off dims = flat_slice+          flat_slice' = FlatSlice off (FlatDimIndex num_segments n : dims)+      out_flat_updated <-+        certifying (distCerts inps aux env) . letExp "flat_update_reg" . BasicOp $+          FlatUpdate as_flat flat_slice' v_flat+      out_updated <-+        letExp "flat_update_reg_reshaped" $+          BasicOp $+            Reshape out_flat_updated $+              reshapeAll as_flat_shape as_lift_shape+      pure $ insertRegulars [distResTag res] [out_updated] env+  | Just as_t <- distInputType <$> lookup as inps,+    isRegularDistResult res,+    not (any (isVariant inps) (flatSliceDims flat_slice)) = do+      -- as should be 1D+      let as_lift_shape = segmentsShape segments <> arrayShape as_t+      as' <- liftSubExpRegular lvl segments inps env as_lift_shape (Var as)+      let update_dims = segments <> flatSliceDims flat_slice+      updated <-+        certifying (distCerts inps aux env)+          . letExp "flat_update_reg_scatter"+          <=< renameExp+          <=< genScatterND lvl as' update_dims+          $ \is -> do+            let (seg_is, in_is) = splitAt (segmentsRank segments) (toList is)+            readInputs segments env seg_is $ filter ((/= as) . fst) inps+            let flat_slice'@(FlatSlice flat_offset _) = fmap pe64 flat_slice+                flat_i =+                  flat_offset + sum (zipWith (*) (map pe64 in_is) (flatSliceStrides flat_slice'))+            -- Value to write+            v' <- letSubExp "v" =<< eIndex v (map toExp in_is)+            -- Index to write `v'` at+            i <- letSubExp "i" =<< toExp flat_i+            pure (seg_is <> [i], v')+      pure $ insertRegulars [distResTag res] [updated] env+  | Just _ <- distInputType <$> lookup as inps = do+      num_segments <- letSubExp "num_segments" =<< toExp (segmentCount segments)+      ns <- letExp "slice_sizes"+        <=< renameExp+        <=< segMap lvl (MkSolo num_segments)+        $ \(MkSolo seg_i) -> do+          seg_is <- segmentCoordsFromFlat segments seg_i+          readInputs segments env seg_is $+            filter ((`elem` flatSliceDims flat_slice) . Var . fst) inps+          slice_dims <- mapM (readInput segments env seg_is inps) $ flatSliceDims flat_slice+          n <- letSubExp "n" <=< toExp $ product $ map pe64 slice_dims+          pure [subExpRes n]+      -- Irregular representation of `as`+      as_rep <- getIrregRep lvl segments env inps as+      IrregularRep shape flags offsets elems _ <-+        ensureDenseIrregular lvl (baseName as <> "_update") as_rep+      -- Inner indices (1 and 2) of `ns`+      (_, _, ii1_vss) <- doRepIota lvl ns+      (_, _, ii2_vss) <- certifying (distCerts inps aux env) $ doSegIota lvl ns+      -- Number of updates to perform+      m <- arraySize 0 <$> lookupType ii2_vss+      elems' <- letExp "elems_scatter" <=< renameExp <=< genScatter lvl elems m $ \gid -> do+        seg_i <- letSubExp "seg_i" =<< eIndex ii1_vss [eSubExp gid]+        in_seg_i <- letSubExp "in_seg_i" =<< eIndex ii2_vss [eSubExp gid]+        seg_is <- segmentCoordsFromFlat segments seg_i+        readInputs segments env seg_is $ filter ((/= as) . fst) inps+        let slice_dims = flatSliceDims flat_slice+            flat_stride = flatSliceStrides flat_slice+            (FlatSlice flat_offset _) = fmap pe64 flat_slice+            in_seg_is =+              unflattenIndex (map pe64 slice_dims) (pe64 in_seg_i)+            flat_i = flat_offset + sum (zipWith (*) in_seg_is (map pe64 flat_stride))+        -- Value to write+        v' <- letSubExp "v" =<< eIndex v (map toExp in_seg_is)+        o' <- letSubExp "o" =<< eIndex offsets [eSubExp seg_i]+        -- Index to write `v'` at+        i <- letExp "i" =<< toExp (pe64 o' + flat_i)+        pure (i, v')+      insertIrregularM shape flags offsets (distResTag res) elems' Dense env+  | otherwise =+      error "Flattening update: destination is not input."++transformRearrange ::+  TrCtx ->+  VName ->+  [Int] ->+  FlattenM DistEnv+transformRearrange (TrCtx lvl segments env inps res aux) v perm+  | isRegularDistResult res = do+      t <- lookupInputType inps v+      v_lifted <-+        liftSubExpRegular+          lvl+          segments+          inps+          env+          (segmentsShape segments <> arrayShape t)+          (Var v)+      let segment_rank = segmentsRank segments+      v_rearrange <- letExp (baseName v <> "_tr") . BasicOp $ Rearrange v_lifted ([0 .. segment_rank - 1] ++ map (+ segment_rank) perm)+      pure $ insertRegulars [distResTag res] [v_rearrange] env+  | otherwise = do+      irreg <- getIrregRep lvl segments env inps v+      t <- lookupInputType inps v+      rep' <-+        certifying (distCerts inps aux env) $+          rearrangeIrreg lvl segments env inps t perm irreg+      insertRepM (distResTag res) (Irregular rep') env++transformScratch ::+  TrCtx ->+  PrimType ->+  [SubExp] ->+  FlattenM DistEnv+transformScratch (TrCtx lvl segments env inps res _aux) pt dims+  | not $ any (isVariant inps) dims = do+      -- All dims are invariant result is regular across segments.+      v' <-+        letExp "scratch" . BasicOp $+          Scratch pt (shapeDims (segmentsShape segments) ++ dims)+      pure $ insertRegulars [distResTag res] [v'] env+  | [n] <- dims = do+      ns <- dataArr lvl segments env inps n+      (_n, offsets, m) <- exScanAndSum lvl ns+      flags <- genFlags lvl m offsets+      res_D <- letExp "scratch_D" $ BasicOp $ Scratch pt [m]+      insertIrregularM ns flags offsets (distResTag res) res_D Dense env+  | otherwise = do+      dim_arrs <- mapM (dataArr lvl segments env inps) dims+      w <- arraySize 0 <$> lookupType (head dim_arrs)+      ns <- letExp "scratch_sizes" <=< segMap lvl (MkSolo w) $ \(MkSolo i) -> do+        vals <- mapM (\arr -> letSubExp "d" =<< eIndex arr [eSubExp i]) dim_arrs+        n <- letSubExp "n" <=< toExp $ product $ map pe64 vals+        pure [subExpRes n]+      (_n, offsets, m) <- exScanAndSum lvl ns+      flags <- genFlags lvl m offsets+      res_D <- letExp "scratch_D" $ BasicOp $ Scratch pt [m]+      insertIrregularM ns flags offsets (distResTag res) res_D Dense env++flattenBasicOp ::+  FlattenOps ->+  Segments ->+  DistEnv ->+  ( DistInputs,+    DistResult,+    PatElem Type,+    StmAux (),+    BasicOp+  ) ->+  FlattenM DistEnv+flattenBasicOp ops segments env (inps, res, pe, aux, e) =+  case e of+    BinOp {} -> scalarCase+    CmpOp {} -> scalarCase+    ConvOp {} -> scalarCase+    UnOp {} -> scalarCase+    UserParam _ _ -> scalarCase -- These are always of type i64.+    Assert {} -> scalarCase+    ArrayLit vs row_type -> transformArrayLit ctx vs row_type+    ArrayVal vs row_type -> transformArrayVal ctx vs row_type+    Opaque op se -> passThrough (Opaque op) se+    Reshape arr reshape -> transformReshape ctx arr reshape+    Index arr slice -> transformIndex ctx arr slice+    FlatIndex arr flat_slice -> transformFlatIndex ctx arr flat_slice+    Iota n x s it -> transformIota ctx n x s it+    Concat d arr shp -> transformConcat ctx d arr shp+    Replicate shape se -> transformReplicate ctx shape se+    Manifest v perm -> transformManifest ctx v perm+    Update safety as slice se -> transformUpdate ctx safety as slice se+    FlatUpdate as flat_slice v -> transformFlatUpdate ctx as flat_slice v+    Rearrange v perm -> transformRearrange ctx v perm+    Scratch pt dims -> transformScratch ctx pt dims+    UpdateAcc {} ->+      -- TODO: handle nonuniform case, which is however rare, and also needs+      -- modifications to WithAcc. The only irregularity that is possible is in+      -- the values to be written.+      scalarCase+    SubExp se -> passThrough SubExp se+  where+    ctx = TrCtx lvl segments env inps res aux+    lvl = flattenSegLevel ops+    scalarCase =+      flattenScalarStm ops segments env inps [res] $+        Let (Pat [pe]) aux (BasicOp e)++    -- Distribute a BasicOp that merely passes its operand through, applying it+    -- to the representation of a distributed input (or falling back to the+    -- scalar case).+    passThrough mkOp se+      | Var v <- se,+        Just (DistInput rt_in _) <- lookup v inps =+          case resVar rt_in env of+            Regular arr -> do+              arr' <-+                certifying (distCerts inps aux env) . letExp (baseName v) $+                  BasicOp (mkOp (Var arr))+              pure $ insertRegulars [distResTag res] [arr'] env+            Irregular irreg -> do+              elems' <-+                certifying (distCerts inps aux env) . letExp (baseName v) $+                  BasicOp (mkOp (Var (irregularD irreg)))+              insertRepM (distResTag res) (Irregular irreg {irregularD = elems'}) env+      | otherwise =+          scalarCase
+ src/Futhark/Pass/Flatten/Builtins.hs view
@@ -0,0 +1,1030 @@+{-# LANGUAGE TypeFamilies #-}++module Futhark.Pass.Flatten.Builtins+  ( BuiltinFn (..),+    builtinFunDef,+    determineReduceOp,+    genUniformSegHist,+    mkSegSpace,+    segMap,+    genFlags,+    genScan,+    genFilter,+    genSegScan,+    genSegScanomap,+    genSegScanomapWithPost,+    genNonSegRed,+    genUniformSegScanomapWithPost,+    genUniformSegRed,+    genSegRed,+    genSegRedomap,+    genScatter,+    genScatterND,+    genShapeIota,+    exScanAndSum,+    genExPrefixSum,+    genSegPrefixSum,+    doRepIota,+    doSegIota,+    doPrefixSum,+    doPartition,+  )+where++import Control.Monad (forM, forM_, (<=<))+import Control.Monad.State.Strict+import Data.Foldable (toList)+import Data.Maybe (fromMaybe)+import Futhark.IR.GPU+import Futhark.IR.SOACS as SOACS+import Futhark.MonadFreshNames+import Futhark.Pass.Flatten.Monad+import Futhark.Tools+import Futhark.Transform.Rename (renameBody, renameLambda)+import Futhark.Transform.ToGPU (getSize, soacsLambdaToGPU)+import Futhark.Util (unsnoc)++mkSegSpace :: (MonadFreshNames m) => [(VName, SubExp)] -> m SegSpace+mkSegSpace dims = SegSpace <$> newVName "phys_tid" <*> pure dims++segIotaName, repIotaName, prefixSumName, partitionName :: Name+segIotaName = builtinName "segiota"+repIotaName = builtinName "repiota"+prefixSumName = builtinName "prefixsum"+partitionName = builtinName "partition"++inlineBuiltinAtLevel :: SegLevel -> Bool+inlineBuiltinAtLevel SegThreadInBlock {} = True+inlineBuiltinAtLevel _ = False++topSegLevel :: SegLevel+topSegLevel = SegThread SegVirt Nothing++data ThreadRecommendation = ManyThreads | NoRecommendation SegVirt++numberOfBlocks ::+  (MonadBuilder m, Op (Rep m) ~ HostOp inner (Rep m)) =>+  Name ->+  SubExp ->+  SubExp ->+  m (SubExp, SubExp)+numberOfBlocks desc w tblock_size = do+  max_num_tblocks_key <- nameFromText . prettyText <$> newVName (desc <> "_num_tblocks")+  num_tblocks <-+    letSubExp "num_tblocks" $+      Op $+        SizeOp $+          CalcNumBlocks w max_num_tblocks_key tblock_size+  num_threads <-+    letSubExp "num_threads" $+      BasicOp $+        BinOp (Mul Int64 OverflowUndef) num_tblocks tblock_size+  pure (num_tblocks, num_threads)++-- | Like 'segThread', but cap the thread count to the input size.+-- This is more efficient for small kernels, e.g. summing a small+-- array.+segThreadCapped ::+  (MonadBuilder m, Rep m ~ GPU) =>+  [SubExp] -> Name -> ThreadRecommendation -> m (SegOpLevel (Rep m))+segThreadCapped ws desc r = do+  w <-+    letSubExp "nest_size"+      =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) ws+  tblock_size <- getSize (desc <> "_tblock_size") SizeThreadBlock++  case r of+    ManyThreads -> do+      usable_groups <-+        letSubExp "segmap_usable_groups"+          =<< eBinOp+            (SDivUp Int64 Unsafe)+            (eSubExp w)+            (eSubExp =<< asIntS Int64 tblock_size)+      let grid = KernelGrid (Count usable_groups) (Count tblock_size)+      pure $ SegThread SegNoVirt (Just grid)+    NoRecommendation v -> do+      (num_tblocks, _) <- numberOfBlocks desc w tblock_size+      let grid = KernelGrid (Count num_tblocks) (Count tblock_size)+      pure $ SegThread v (Just grid)++-- FIXME: We use segThreadCapped here because otherwise we may get+-- out-of-bounds writes for SegOps with non-primitive return types.+capThreadSegLevel ::+  (MonadBuilder m, Rep m ~ GPU, Foldable t) =>+  t SubExp -> Name -> SegLevel -> ThreadRecommendation -> m SegLevel+capThreadSegLevel segments desc lvl tr =+  case lvl of+    SegThread {} -> segThreadCapped (toList segments) desc tr+    _ -> pure lvl++determineReduceOp ::+  (MonadBuilder m) =>+  Lambda SOACS ->+  [SubExp] ->+  m (Lambda SOACS, [SubExp], Shape)+determineReduceOp lam nes =+  -- We obtain the scalar neutral element by indexing the array-typed+  -- one at [0,...,0]. This is safe even if the array is not literally+  -- a replicate: every lane of a vectorised operator must have a+  -- neutral element, and neutral elements are unique, so all lanes+  -- are forced to hold the same value.+  case mapM subExpVar nes of+    Just ne_vs' -> do+      let (shape, lam') = isVectorMap lam+      nes' <- forM ne_vs' $ \ne_v -> do+        ne_v_t <- lookupType ne_v+        letSubExp "hist_ne" $+          BasicOp $+            Index ne_v $+              fullSlice ne_v_t $+                replicate (shapeRank shape) $+                  DimFix $+                    intConst Int64 0+      pure (lam', nes', shape)+    Nothing ->+      pure (lam, nes, mempty)++isVectorMap :: Lambda SOACS -> (Shape, Lambda SOACS)+isVectorMap lam+  | [Let (Pat pes) _ (Op (Screma w arrs form))] <-+      stmsToList $ bodyStms $ lambdaBody lam,+    map resSubExp (bodyResult (lambdaBody lam)) == map (Var . patElemName) pes,+    Just map_lam <- isMapSOAC form,+    arrs == map paramName (lambdaParams lam) =+      let (shape, lam') = isVectorMap map_lam+       in (Shape [w] <> shape, lam')+  | otherwise = (mempty, lam)++segMap ::+  (MonadBuilder m, Rep m ~ GPU, Traversable f) =>+  SegLevel -> f SubExp -> (f SubExp -> m Result) -> m (Exp GPU)+segMap lvl segments f = do+  gtids <- traverse (const $ newVName "gtid") segments+  space <- mkSegSpace $ zip (toList gtids) (toList segments)+  ((res, ts), stms) <- collectStms $ localScope (scopeOfSegSpace space) $ do+    res <- f $ fmap Var gtids+    ts <- mapM (subExpType . resSubExp) res+    pure (map mkResult res, ts)+  let kbody = Body () stms res+  let tr = if all primType ts then ManyThreads else NoRecommendation SegVirt+  lvl' <- capThreadSegLevel segments "segmap" lvl tr+  pure $ Op $ SegOp $ SegMap lvl' space ts kbody+  where+    mkResult (SubExpRes cs se) = Returns ResultMaySimplify cs se++genScanWithKernelBody ::+  (MonadBuilder m, Rep m ~ GPU, Traversable f) =>+  SegLevel ->+  Name ->+  f SubExp ->+  Lambda GPU ->+  [SubExp] ->+  (f SubExp -> m Result) ->+  m [VName]+genScanWithKernelBody lvl desc segments lam nes =+  genScanWithKernelBodyAndPost+    lvl+    desc+    segments+    (\_ -> pure lam)+    mempty+    nes+    (\_ res_t -> mkIdentityLambda res_t)++-- The SegSpace of a SegRed must always have at least two dimensions, so that+-- its result is an array. We therefore prepend a dummy unit dimension, and+-- afterwards index out the single element of the unit-sized result.+genNonSegRed ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel ->+  Name ->+  [SubExp] ->+  Reduce GPU ->+  Shape ->+  Lambda GPU ->+  [VName] ->+  m [VName]+genNonSegRed lvl desc segments red_op shape map_lam arrs = do+  let red_lam = redLambda red_op+      nes = redNeutral red_op+      comm = redComm red_op+  let dummy = intConst Int64 1+  gtids_dummy <- newVName "dummy"+  gtids_original <- traverse (const $ newVName "gtid") segments+  let gtids = gtids_dummy : gtids_original+  let new_segment = dummy : segments+  space <- mkSegSpace $ zip (toList gtids) (toList new_segment)+  let gtids' = fmap Var gtids+  (kbody, res_t) <- buildBody . localScope (scopeOfSegSpace space) $ do+    bindLambdaInputArrays (drop 1 gtids') map_lam arrs+    res <- bodyBind (lambdaBody map_lam)+    res_t <- mapM (subExpType . resSubExp) res+    pure (map mkResult res, res_t)+  kbody' <- renameBody kbody+  red_lam' <- renameLambda red_lam+  let op = SegBinOp comm red_lam' nes shape+  lvl' <- capThreadSegLevel new_segment "uniform_nonsegred" lvl $ NoRecommendation SegNoVirt+  ress <- letTupExp desc $ Op $ SegOp $ SegRed lvl' space res_t kbody' [op]+  forM ress $ \res_d -> do+    res_dt <- lookupType res_d+    letExp desc . BasicOp $+      case res_dt of+        Acc {} -> SubExp $ Var res_d+        _ -> Index res_d $ fullSlice res_dt [DimFix $ intConst Int64 0]+  where+    mkResult (SubExpRes cs se) = Returns ResultMaySimplify cs se++genUniformSegHist ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel ->+  Name ->+  [SubExp] ->+  [SOACS.HistOp SOACS] ->+  Lambda GPU ->+  [VName] ->+  ([SubExp] -> m ()) ->+  m [VName]+genUniformSegHist lvl desc segments ops bucket_fun arrs readFree = do+  ops' <- forM ops $ \(SOACS.HistOp dest_shape rf dests nes op) -> do+    (op', nes', shape) <- determineReduceOp op nes+    let op'' = soacsLambdaToGPU op'+    pure $ Futhark.IR.GPU.HistOp dest_shape rf dests nes' shape op''+  gtids <- traverse (const $ newVName "gtid") segments+  space <- mkSegSpace $ zip (toList gtids) (toList segments)+  let gtids' = fmap Var gtids+  (kbody, res_t) <- buildBody . localScope (scopeOfSegSpace space) $ do+    readFree gtids'+    bindLambdaInputArrays gtids' bucket_fun arrs+    res <- bodyBind (lambdaBody bucket_fun)+    res_t <- mapM (subExpType . resSubExp) res+    pure (map mkResult res, res_t)+  kbody' <- renameBody kbody+  lvl' <- capThreadSegLevel segments "uniform_seghist" lvl $ NoRecommendation SegNoVirt+  letTupExp desc $ Op $ SegOp $ SegHist lvl' space res_t kbody' ops'+  where+    mkResult (SubExpRes cs se) = Returns ResultMaySimplify cs se++genUniformSegRed ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel ->+  Name ->+  [SubExp] ->+  [Reduce GPU] ->+  Shape ->+  Lambda GPU ->+  [VName] ->+  ([SubExp] -> m ()) ->+  m [VName]+genUniformSegRed lvl desc segments red_ops shape map_lam arrs readFree = do+  gtids <- traverse (const $ newVName "gtid") segments+  space <- mkSegSpace $ zip (toList gtids) (toList segments)+  let gtids' = fmap Var gtids+  (kbody, res_t) <- buildBody . localScope (scopeOfSegSpace space) $ do+    readFree gtids'+    bindLambdaInputArrays gtids' map_lam arrs+    res <- bodyBind (lambdaBody map_lam)+    res_t <- mapM (subExpType . resSubExp) res+    pure (map mkResult res, res_t)++  ops <- forM red_ops $ \red_op -> do+    red_lam' <- renameLambda $ redLambda red_op+    pure $ SegBinOp (redComm red_op) red_lam' (redNeutral red_op) shape++  kbody' <- renameBody kbody+  lvl' <- capThreadSegLevel segments "uniform_segred" lvl $ NoRecommendation SegNoVirt+  letTupExp desc $ Op $ SegOp $ SegRed lvl' space res_t kbody' ops+  where+    mkResult (SubExpRes cs se) = Returns ResultMaySimplify cs se++genScanWithKernelBodyAndPost ::+  (MonadBuilder m, Rep m ~ GPU, Traversable f) =>+  SegLevel ->+  Name ->+  f SubExp ->+  (f SubExp -> m (Lambda GPU)) ->+  Shape ->+  [SubExp] ->+  (f SubExp -> [Type] -> m (Lambda GPU)) ->+  (f SubExp -> m Result) ->+  m [VName]+genScanWithKernelBodyAndPost lvl desc segments mkScanLam shape nes mkPostLam m = do+  gtids <- traverse (const $ newVName "gtid") segments+  space <- mkSegSpace $ zip (toList gtids) (toList segments)+  let gtids' = fmap Var gtids+  (kbody, res_t) <- buildBody . localScope (scopeOfSegSpace space) $ do+    res <- m gtids'+    res_t <- mapM (subExpType . resSubExp) res+    pure (map mkResult res, res_t)++  scan_lam <- mkScanLam gtids'+  post_lam <-+    localScope (scopeOfSegSpace space) $+      mkPostLam gtids' res_t+  -- We have to rename since we are using a global readFree+  scan_lam' <- renameLambda scan_lam+  post_lam' <- renameLambda post_lam++  kbody' <- renameBody kbody+  let op = SegBinOp Noncommutative scan_lam' nes shape+  lvl' <- capThreadSegLevel segments "uniform_segscan" lvl $ NoRecommendation SegNoVirt+  letTupExp desc $ Op $ SegOp $ SegScan lvl' space res_t kbody' [op] (SegPostOp post_lam')+  where+    mkResult (SubExpRes cs se) = Returns ResultMaySimplify cs se++bindLambdaInputArrays ::+  (MonadBuilder m, Traversable f) =>+  f SubExp ->+  Lambda GPU ->+  [VName] ->+  m ()+bindLambdaInputArrays gtids lam arrs = do+  let idxs = toList gtids+  forM_ (zip (lambdaParams lam) arrs) $ \(p, arr) ->+    letBindNames [paramName p]+      =<< case paramType p of+        Acc {} ->+          eSubExp $ Var arr+        _ ->+          eIndex arr $ map eSubExp idxs++genScan ::+  (Traversable f, MonadBuilder m, Rep m ~ GPU) =>+  SegLevel -> Name -> f SubExp -> Lambda GPU -> [SubExp] -> [VName] -> m [VName]+genScan lvl desc segments lam nes arrs =+  genScanWithKernelBody lvl desc segments lam nes $ \gtids ->+    fmap subExpsRes . forM arrs $ \arr ->+      letSubExp (baseName arr <> "_elem") =<< eIndex arr (toList $ fmap eSubExp gtids)++genExScan ::+  (MonadBuilder m, Rep m ~ GPU, Traversable f) =>+  SegLevel -> Name -> f SubExp -> Lambda GPU -> [SubExp] -> [VName] -> m [VName]+genExScan lvl desc segments lam nes arrs =+  genScanWithKernelBody lvl desc segments lam nes $ \gtids ->+    let Just (outerDims, innerDim) = unsnoc $ toList gtids+     in do+          prescan <-+            letTupExp' "to_prescan"+              =<< eIf+                (toExp $ pe64 innerDim .==. 0)+                (eBody (map eSubExp nes))+                (eBody (map (`eIndex` (map toExp outerDims ++ [toExp $ pe64 innerDim - 1])) arrs))+          pure $ subExpsRes prescan++segScanLambda ::+  (MonadBuilder m, Rep m ~ GPU) =>+  Lambda GPU ->+  ([SubExp] -> m ()) ->+  [SubExp] ->+  m (Lambda GPU)+segScanLambda lam _readFree _gtids = do+  x_flag_p <- newParam "x_flag" $ Prim Bool+  y_flag_p <- newParam "y_flag" $ Prim Bool+  let ts = lambdaReturnType lam+      (xps, yps) = splitAt (length ts) $ lambdaParams lam+  mkLambda ([x_flag_p] ++ xps ++ [y_flag_p] ++ yps) $+    bodyBind+      =<< eBody+        [ eBinOp LogOr (eParam x_flag_p) (eParam y_flag_p),+          eIf+            (eParam y_flag_p)+            (eBody (map eParam yps))+            (pure $ lambdaBody lam)+        ]++genSegScan ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel -> Name -> Lambda GPU -> [SubExp] -> VName -> [VName] -> m [VName]+genSegScan lvl desc lam nes flags arrs = do+  w <- arraySize 0 <$> lookupType flags+  lam' <- segScanLambda lam (const $ pure ()) []+  drop 1 <$> genScan lvl desc [w] lam' (constant False : nes) (flags : arrs)++segScanomapPostLambda ::+  (MonadBuilder m, Rep m ~ GPU) =>+  Lambda GPU ->+  ([SubExp] -> m ()) ->+  [SubExp] ->+  m (Lambda GPU)+segScanomapPostLambda lam readFree gtids = do+  flag_p <- newParam "seg_flag" $ Prim Bool+  mkLambda (flag_p : lambdaParams lam) $ do+    readFree gtids+    bodyBind $ lambdaBody lam++genSegScanomap ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel ->+  Name ->+  Lambda GPU ->+  [SubExp] ->+  VName ->+  Lambda GPU ->+  [VName] ->+  ([SubExp] -> m ()) ->+  m [VName]+genSegScanomap lvl desc scan_lam nes flags map_lam arrs readFree = do+  post_lam <- mkIdentityLambda $ lambdaReturnType map_lam+  genSegScanomapWithPost lvl desc scan_lam nes flags post_lam map_lam arrs readFree++genSegScanomapWithPost ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel ->+  Name ->+  Lambda GPU ->+  [SubExp] ->+  VName ->+  Lambda GPU ->+  Lambda GPU ->+  [VName] ->+  ([SubExp] -> m ()) ->+  m [VName]+genSegScanomapWithPost lvl desc scan_lam nes flags post_lam map_lam arrs readFree = do+  w <- arraySize 0 <$> lookupType flags++  genScanWithKernelBodyAndPost+    lvl+    desc+    [w]+    (segScanLambda scan_lam readFree)+    mempty+    (constant False : nes)+    ( \gtids _res_t ->+        segScanomapPostLambda post_lam readFree gtids+    )+    ( \gtids -> do+        let [gtid] = toList gtids+        flag <- letSubExp "flag" =<< eIndex flags [eSubExp gtid]+        readFree gtids+        bindLambdaInputArrays gtids map_lam arrs+        map_res <- bodyBind (lambdaBody map_lam)+        pure (subExpRes flag : map_res)+    )++withReadFree ::+  (MonadBuilder m, Rep m ~ GPU) =>+  Lambda GPU ->+  ([SubExp] -> m ()) ->+  [SubExp] ->+  m (Lambda GPU)+withReadFree lam readFree gtids =+  mkLambda (lambdaParams lam) $ do+    readFree gtids+    bodyBind $ lambdaBody lam++genUniformSegScanomapWithPost ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel ->+  [SubExp] ->+  Name ->+  Lambda GPU ->+  Shape ->+  [SubExp] ->+  Lambda GPU ->+  Lambda GPU ->+  [VName] ->+  ([SubExp] -> m ()) ->+  m [VName]+genUniformSegScanomapWithPost lvl segments desc scan_lam shape nes post_lam map_lam arrs readFree = do+  genScanWithKernelBodyAndPost+    lvl+    desc+    segments+    (const $ pure scan_lam)+    shape+    nes+    (\gtids _res_t -> withReadFree post_lam readFree gtids)+    ( \gtids -> do+        readFree gtids+        bindLambdaInputArrays gtids map_lam arrs+        bodyBind (lambdaBody map_lam)+    )++genPrefixSum ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel -> Name -> VName -> m VName+genPrefixSum lvl desc ns = do+  ws <- arrayDims <$> lookupType ns+  add_lam <- binOpLambda (Add Int64 OverflowUndef) int64+  head <$> genScan lvl desc ws add_lam [intConst Int64 0] [ns]++genExPrefixSum ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel -> Name -> VName -> m VName+genExPrefixSum lvl desc ns = do+  ws <- arrayDims <$> lookupType ns+  add_lam <- binOpLambda (Add Int64 OverflowUndef) int64+  head <$> genExScan lvl desc ws add_lam [intConst Int64 0] [ns]++genSegPrefixSum ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel -> Name -> VName -> VName -> m VName+genSegPrefixSum lvl desc flags ns = do+  add_lam <- binOpLambda (Add Int64 OverflowUndef) int64+  head <$> genSegScan lvl desc add_lam [intConst Int64 0] flags [ns]++-- | Convenience wrapper around 'genScatterND' for one-dimensional+-- destinations.+genScatter ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel -> VName -> SubExp -> (SubExp -> m (VName, SubExp)) -> m (Exp GPU)+genScatter lvl dest n f =+  genScatterND lvl dest [n] $ \ ~[gtid] -> do+    (i, v) <- f gtid+    pure ([Var i], v)++genScatterND ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel -> VName -> [SubExp] -> ([SubExp] -> m ([SubExp], SubExp)) -> m (Exp GPU)+genScatterND lvl dest grid f = do+  gtids <- traverse (const $ newVName "gtid") grid+  space <- mkSegSpace $ zip gtids grid+  dest_t <- lookupType dest+  let accRank = arrayRank dest_t+  withAcc [dest] accRank $ \ ~[acc] -> do+    kbody <- buildBody_ $ localScope (scopeOfSegSpace space) $ do+      (idxs, v) <- f $ fmap Var gtids+      acc' <-+        letExp (baseName acc) $+          BasicOp $+            UpdateAcc Safe acc idxs [v]+      pure [Returns ResultMaySimplify mempty $ Var acc']+    acc_t <- lookupType acc+    lvl' <- capThreadSegLevel grid "genScatterND" lvl $ NoRecommendation SegVirt+    letTupExp' "scatter" $ Op $ SegOp $ SegMap lvl' space [acc_t] kbody++genTabulate ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel -> SubExp -> (SubExp -> m [SubExp]) -> m (Exp GPU)+genTabulate lvl w m = do+  gtid <- newVName "gtid"+  space <- mkSegSpace [(gtid, w)]+  ((res, ts), stms) <- collectStms $ localScope (scopeOfSegSpace space) $ do+    ses <- m $ Var gtid+    ts <- mapM subExpType ses+    pure (map (Returns ResultMaySimplify mempty) ses, ts)+  let kbody = Body () stms res+  lvl' <- capThreadSegLevel [w] "genTabulate" lvl $ NoRecommendation SegVirt+  pure $ Op $ SegOp $ SegMap lvl' space ts kbody++genFlags ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel -> SubExp -> VName -> m VName+genFlags lvl m offsets = do+  flags_allfalse <-+    letExp "flags_allfalse" . BasicOp $+      Replicate (Shape [m]) (constant False)+  n <- arraySize 0 <$> lookupType offsets+  letExp "flags" <=< genScatter lvl flags_allfalse n $ \gtid -> do+    i <- letExp "i" =<< eIndex offsets [eSubExp gtid]+    pure (i, constant True)++genSegRed ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel -> VName -> VName -> VName -> [VName] -> Reduce SOACS -> m [VName]+genSegRed lvl segments flags offsets elems red = do+  scanned <-+    genSegScan+      lvl+      "red"+      (soacsLambdaToGPU $ redLambda red)+      (redNeutral red)+      flags+      elems+  num_segments <- arraySize 0 <$> lookupType offsets+  letTupExp "segred" <=< genTabulate lvl num_segments $ \i -> do+    n <- letSubExp "n" =<< eIndex segments [eSubExp i]+    offset <- letSubExp "offset" =<< eIndex offsets [toExp (pe64 i)]+    letTupExp' "segment_res" <=< eIf (toExp $ pe64 n .==. 0) (eBody $ map eSubExp nes) $+      eBody $+        map (`eIndex` [toExp $ pe64 offset + pe64 n - 1]) scanned+  where+    nes = redNeutral red++genSegRedomap ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel ->+  VName ->+  VName ->+  VName ->+  [VName] ->+  Reduce SOACS ->+  Lambda GPU ->+  ([SubExp] -> m ()) ->+  m ([VName], [VName])+genSegRedomap lvl segments flags offsets elems red map_lam readFree = do+  scanned_and_map <-+    genSegScanomap+      lvl+      "redomap"+      (soacsLambdaToGPU $ redLambda red)+      (redNeutral red)+      flags+      map_lam+      elems+      readFree+  let (scanned, mapout) = splitAt (length nes) scanned_and_map+  num_segments <- arraySize 0 <$> lookupType offsets+  reds <- letTupExp "segred" <=< genTabulate lvl num_segments $ \i -> do+    n <- letSubExp "n" =<< eIndex segments [eSubExp i]+    offset <- letSubExp "offset" =<< eIndex offsets [toExp (pe64 i)]+    letTupExp' "segment_res" <=< eIf (toExp $ pe64 n .==. 0) (eBody $ map eSubExp nes) $+      eBody $+        map (`eIndex` [toExp $ pe64 offset + pe64 n - 1]) scanned+  pure (reds, mapout)+  where+    nes = redNeutral red++-- | Produces a multidimensional iota for the given shape.+genShapeIota ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel -> Shape -> m VName+genShapeIota lvl shape = do+  let dims = shapeDims shape+  letExp "shape_iota" <=< segMap lvl dims $ \gtids -> do+    i <-+      toSubExp "shape_iota_elem" $+        flattenIndex (map pe64 dims) (map pe64 gtids)+    pure [subExpRes i]++-- Returns (#segments, segment start offsets, sum of segment sizes)+-- Note: If given a multi-dimensional array,+-- `#segments` and `sum of segment sizes` will be arrays, not scalars.+-- `segment start offsets` will always have the same shape as `ks`.+exScanAndSum ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel -> VName -> m (SubExp, VName, SubExp)+exScanAndSum lvl ks = do+  ns <- arrayDims <$> lookupType ks+  -- If `ks` only has a single dimension+  -- the size will be a scalar, otherwise it's an array.+  ns' <- letExp "ns" $ BasicOp $ case ns of+    [] -> error $ "exScanAndSum: Given non-array argument: " ++ prettyString ks+    [n] -> SubExp n+    _ -> ArrayLit ns (Prim int64)+  -- Check if the innermost dimension is empty.+  is_empty <-+    letExp "is_empty"+      =<< ( case ns of+              [n] -> toExp (pe64 n .==. 0)+              _ -> eLast ns' >>= letSubExp "n" >>= (\n -> toExp $ pe64 n .==. 0)+          )+  offsets <- letExp "offsets" =<< toExp =<< genExPrefixSum lvl "offsets" ks+  ms <- letExp "ms" <=< segMap lvl (init ns) $ \gtids -> do+    let idxs = map toExp gtids+    offset <- letExp "offset" =<< eIndex offsets idxs+    k <- letExp "k" =<< eIndex ks idxs+    m <-+      letSubExp "m"+        =<< eIf+          (toExp is_empty)+          (eBody [eSubExp $ intConst Int64 0])+          -- Add last size because 'offsets' is an *exclusive* prefix+          -- sum.+          (eBody [eBinOp (Add Int64 OverflowUndef) (eLast offset) (eLast k)])+    pure [subExpRes m]+  pure (Var ns', offsets, Var ms)++genSegIota ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel -> VName -> m (VName, VName, VName)+genSegIota lvl ks = do+  (_n, offsets, m) <- exScanAndSum lvl ks+  flags <- genFlags lvl m offsets+  ones <- letExp "ones" $ BasicOp $ Replicate (Shape [m]) one+  iotas <- genSegPrefixSum lvl "iotas" flags ones+  res <- letExp "res" <=< genTabulate lvl m $ \i -> do+    x <- letSubExp "x" =<< eIndex iotas [eSubExp i]+    letTupExp' "xm1" $ BasicOp $ BinOp (Sub Int64 OverflowUndef) x one+  pure (flags, offsets, res)+  where+    one = intConst Int64 1++genRepIota ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel ->+  VName ->+  m (VName, VName, VName)+genRepIota lvl ks = do+  (n, offsets, m) <- exScanAndSum lvl ks+  is <- letExp "is" <=< genTabulate lvl n $ \i -> do+    o <- letSubExp "o" =<< eIndex offsets [eSubExp i]+    k <- letSubExp "n" =<< eIndex ks [eSubExp i]+    letTupExp' "i"+      =<< eIf+        (toExp (pe64 k .==. 0))+        (eBody [eSubExp negone])+        (eBody [toExp $ pe64 o])+  zeroes <- letExp "zeroes" $ BasicOp $ Replicate (Shape [m]) zero+  starts <-+    letExp "starts" <=< genScatter lvl zeroes n $ \gtid -> do+      i <- letExp "i" =<< eIndex is [eSubExp gtid]+      pure (i, gtid)+  flags <- genFlags lvl m offsets+  res <- genSegPrefixSum lvl "res" flags starts+  pure (flags, offsets, res)+  where+    zero = intConst Int64 0+    negone = intConst Int64 (-1)++genPartition ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel -> VName -> VName -> VName -> m (VName, VName, VName)+genPartition lvl n k cls = do+  let n' = Var n+  let k' = Var k+  let dims = [k', n']+  -- Create a `[k][n]` array of flags such that `cls_flags[i][j]`+  -- is equal 1 if the j'th element is a member of equivalence class `i` i.e.+  -- the `i`th row is a flag array for equivalence class `i`.+  cls_flags <-+    letExp "flags"+      <=< segMap lvl dims+      $ \[i, j] -> do+        c <- letSubExp "c" =<< eIndex cls [toExp j]+        cls_flag <-+          letSubExp "cls_flag"+            =<< eIf+              (toExp $ pe64 i .==. pe64 c)+              (eBody [toExp $ intConst Int64 1])+              (eBody [toExp $ intConst Int64 0])+        pure [subExpRes cls_flag]++  -- Offsets of each of the individual equivalence classes.+  (_, local_offs, _counts) <- exScanAndSum lvl cls_flags+  -- The number of elems in each class+  counts <- letExp "counts" =<< toExp _counts+  -- Offsets of the whole equivalence classes+  global_offs <- genExPrefixSum lvl "global_offs" counts+  -- Offsets over all of the equivalence classes.+  cls_offs <-+    letExp "cls_offs" =<< do+      segMap lvl dims $ \[i, j] -> do+        global_offset <- letExp "global_offset" =<< eIndex global_offs [toExp i]+        offset <-+          letSubExp "offset"+            =<< eBinOp+              (Add Int64 OverflowUndef)+              (eIndex local_offs [toExp i, toExp j])+              (toExp global_offset)+        pure [subExpRes offset]++  scratch <- letExp "scratch" $ BasicOp $ Scratch int64 [n']+  res <- letExp "scatter_res" <=< genScatter lvl scratch n' $ \gtid -> do+    c <- letExp "c" =<< eIndex cls [toExp gtid]+    ind <- letExp "ind" =<< eIndex cls_offs [toExp c, toExp gtid]+    i <- letSubExp "i" =<< toExp gtid+    pure (ind, i)+  pure (counts, global_offs, res)++genFilter ::+  (MonadBuilder m, Rep m ~ GPU) =>+  SegLevel -> VName -> m (SubExp, VName)+genFilter lvl flags = do+  w <- arraySize 0 <$> lookupType flags+  flags_int <- letExp "flags_int" <=< segMap lvl [w] $ \[i] -> do+    b <- letSubExp "b" =<< eIndex flags [eSubExp i]+    v <-+      letSubExp "v"+        =<< eIf+          (eSubExp b)+          (eBody [toExp $ intConst Int64 1])+          (eBody [toExp $ intConst Int64 0])+    pure [subExpRes v]+  -- offsets <- genExPrefixSum "filter_offs" flags_int+  (_n, offsets, num_true) <- exScanAndSum lvl flags_int+  -- num_true <- letSubExp "num_true"  =<< eIndex flags_int [toExp $ pe64 w - 1]+  scratch <- letExp "scratch" $ BasicOp $ Scratch int64 [num_true]+  -- is this efficient or do i need to do something smarter? like scatter with guard?+  -- offsets' <- letExp "offset" <=< segMap [w] $ \[i] -> do+  --   b' <- letSubExp "b" =<< eIndex flags [eSubExp i]+  --   v' <-+  --     letSubExp "v'"+  --       =<< eIf+  --         (eSubExp b')+  --         (eBody [eIndex offsets [eSubExp i]] )+  --         (eBody [toExp $ intConst Int64 (-1)])+  --   pure [subExpRes v']++  filtered <- letExp "filtered" <=< genScatter lvl scratch w $ \gtid -> do+    b <- letSubExp "b" =<< eIndex flags [eSubExp gtid]+    -- idx <- letExp "idx" =<< eIndex offsets' [eSubExp gtid]+    idx_se <-+      letSubExp "idx"+        =<< eIf+          (eSubExp b)+          (eBody [eIndex offsets [eSubExp gtid]])+          (eBody [toExp $ intConst Int64 (-1)])+    -- maybe cleaner?+    idx <- letExp "idx" =<< toExp idx_se+    pure (idx, gtid)+  pure (num_true, filtered)++buildingBuiltin :: Builder GPU (FunDef GPU) -> FunDef GPU+buildingBuiltin m = fst $ evalState (runBuilderT m mempty) blankNameSource++segIotaBuiltin :: FunDef GPU+segIotaBuiltin = buildingBuiltin $ do+  np <- newParam "n" $ Prim int64+  nsp <- newParam "ns" $ Array int64 (Shape [Var (paramName np)]) Nonunique+  body <-+    localScope (scopeOfFParams [np, nsp]) . buildBody_ $ do+      (flags, offsets, res) <- genSegIota topSegLevel (paramName nsp)+      m <- arraySize 0 <$> lookupType res+      pure $ subExpsRes [m, Var flags, Var offsets, Var res]+  pure+    FunDef+      { funDefEntryPoint = Nothing,+        funDefAttrs = mempty,+        funDefName = segIotaName,+        funDefRetType =+          map+            (,mempty)+            [ Prim int64,+              Array Bool (Shape [Ext 0]) Unique,+              Array int64 (Shape [Free $ Var $ paramName np]) Unique,+              Array int64 (Shape [Ext 0]) Unique+            ],+        funDefParams = [np, nsp],+        funDefBody = body+      }++repIotaBuiltin :: FunDef GPU+repIotaBuiltin = buildingBuiltin $ do+  np <- newParam "n" $ Prim int64+  nsp <- newParam "ns" $ Array int64 (Shape [Var (paramName np)]) Nonunique+  body <-+    localScope (scopeOfFParams [np, nsp]) . buildBody_ $ do+      (flags, offsets, res) <- genRepIota topSegLevel (paramName nsp)+      m <- arraySize 0 <$> lookupType res+      pure $ subExpsRes [m, Var flags, Var offsets, Var res]+  pure+    FunDef+      { funDefEntryPoint = Nothing,+        funDefAttrs = mempty,+        funDefName = repIotaName,+        funDefRetType =+          map+            (,mempty)+            [ Prim int64,+              Array Bool (Shape [Ext 0]) Unique,+              Array int64 (Shape [Free $ Var $ paramName np]) Unique,+              Array int64 (Shape [Ext 0]) Unique+            ],+        funDefParams = [np, nsp],+        funDefBody = body+      }++prefixSumBuiltin :: FunDef GPU+prefixSumBuiltin = buildingBuiltin $ do+  np <- newParam "n" $ Prim int64+  nsp <- newParam "ns" $ Array int64 (Shape [Var (paramName np)]) Nonunique+  body <-+    localScope (scopeOfFParams [np, nsp]) . buildBody_ $+      varsRes . pure <$> genPrefixSum topSegLevel "res" (paramName nsp)+  pure+    FunDef+      { funDefEntryPoint = Nothing,+        funDefAttrs = mempty,+        funDefName = prefixSumName,+        funDefRetType =+          [(Array int64 (Shape [Free $ Var $ paramName np]) Unique, mempty)],+        funDefParams = [np, nsp],+        funDefBody = body+      }++partitionBuiltin :: FunDef GPU+partitionBuiltin = buildingBuiltin $ do+  np <- newParam "n" $ Prim int64+  kp <- newParam "k" $ Prim int64+  csp <- newParam "cs" $ Array int64 (Shape [Var (paramName np)]) Nonunique+  body <-+    localScope (scopeOfFParams [np, kp, csp]) . buildBody_ $ do+      (counts, offsets, res) <- genPartition topSegLevel (paramName np) (paramName kp) (paramName csp)+      pure $ varsRes [counts, offsets, res]+  pure+    FunDef+      { funDefEntryPoint = Nothing,+        funDefAttrs = mempty,+        funDefName = partitionName,+        funDefRetType =+          map+            (,mempty)+            [ Array int64 (Shape [Free $ Var $ paramName kp]) Unique,+              Array int64 (Shape [Free $ Var $ paramName kp]) Unique,+              Array int64 (Shape [Free $ Var $ paramName np]) Unique+            ],+        funDefParams = [np, kp, csp],+        funDefBody = body+      }++-- | Retrieve the function definition corresponding to a builtin.+builtinFunDef :: BuiltinFn -> FunDef GPU+builtinFunDef BuiltinSegIota = segIotaBuiltin+builtinFunDef BuiltinRepIota = repIotaBuiltin+builtinFunDef BuiltinPrefixSum = prefixSumBuiltin+builtinFunDef BuiltinPartition = partitionBuiltin++-- | @[0,1,2,0,1,0,1,2,3,4,...]@.  Returns @(flags,offsets,elems)@.+doSegIota ::+  SegLevel -> VName -> FlattenM (VName, VName, VName)+doSegIota lvl ns+  | inlineBuiltinAtLevel lvl =+      genSegIota lvl ns+  | otherwise = do+      demandBuiltin BuiltinSegIota+      ns_t <- lookupType ns+      let n = arraySize 0 ns_t+      m <- newVName "m"+      flags <- newVName "segiota_flags"+      offsets <- newVName "segiota_offsets"+      elems <- newVName "segiota_elems"+      let args = [(n, Prim int64), (Var ns, ns_t)]+          restype =+            fromMaybe (error "doSegIota: bad application") $+              applyRetType+                (map fst $ funDefRetType segIotaBuiltin)+                (funDefParams segIotaBuiltin)+                args+      letBindNames [m, flags, offsets, elems] $+        Apply+          (funDefName segIotaBuiltin)+          [(n, Observe), (Var ns, Observe)]+          (map (,mempty) restype)+          Safe+      pure (flags, offsets, elems)++-- | Produces @[0,0,0,1,1,2,2,2,...]@.  Returns @(flags, offsets,+-- elems)@.+doRepIota ::+  SegLevel -> VName -> FlattenM (VName, VName, VName)+doRepIota lvl ns+  | inlineBuiltinAtLevel lvl =+      genRepIota lvl ns+  | otherwise = do+      demandBuiltin BuiltinRepIota+      ns_t <- lookupType ns+      let n = arraySize 0 ns_t+      m <- newVName "m"+      flags <- newVName "repiota_flags"+      offsets <- newVName "repiota_offsets"+      elems <- newVName "repiota_elems"+      let args = [(n, Prim int64), (Var ns, ns_t)]+          restype =+            fromMaybe (error "doRepIota: bad application") $+              applyRetType+                (map fst $ funDefRetType repIotaBuiltin)+                (funDefParams repIotaBuiltin)+                args+      letBindNames [m, flags, offsets, elems] $+        Apply+          (funDefName repIotaBuiltin)+          [(n, Observe), (Var ns, Observe)]+          (map (,mempty) restype)+          Safe+      pure (flags, offsets, elems)++doPrefixSum ::+  SegLevel -> VName -> FlattenM VName+doPrefixSum lvl ns+  | inlineBuiltinAtLevel lvl =+      genPrefixSum lvl "prefix_sum" ns+  | otherwise = do+      demandBuiltin BuiltinPrefixSum+      ns_t <- lookupType ns+      let n = arraySize 0 ns_t+      letExp "prefix_sum" $+        Apply+          (funDefName prefixSumBuiltin)+          [(n, Observe), (Var ns, Observe)]+          [(toDecl (staticShapes1 ns_t) Unique, mempty)]+          Safe++doPartition ::+  SegLevel -> VName -> VName -> FlattenM (VName, VName, VName)+doPartition lvl k cs+  | inlineBuiltinAtLevel lvl = do+      cs_t <- lookupType cs+      n <- letExp "n" $ BasicOp $ SubExp $ arraySize 0 cs_t+      genPartition lvl n k cs+  | otherwise = do+      demandBuiltin BuiltinPartition+      cs_t <- lookupType cs+      let n = arraySize 0 cs_t+      counts <- newVName "partition_counts"+      offsets <- newVName "partition_offsets"+      res <- newVName "partition_res"+      let args = [(n, Prim int64), (Var k, Prim int64), (Var cs, cs_t)]+          restype =+            fromMaybe (error "doPartition: bad application") $+              applyRetType+                (map fst $ funDefRetType partitionBuiltin)+                (funDefParams partitionBuiltin)+                args+      letBindNames [counts, offsets, res] $+        Apply+          (funDefName partitionBuiltin)+          [(n, Observe), (Var k, Observe), (Var cs, Observe)]+          (map (,mempty) restype)+          Safe+      pure (counts, offsets, res)
+ src/Futhark/Pass/Flatten/Distribute.hs view
@@ -0,0 +1,674 @@+module Futhark.Pass.Flatten.Distribute+  ( distributeMap,+    distributeBody,+    MapArray (..),+    mapArrayRowType,+    DistResults (..),+    DistRep,+    ResMap,+    Distributed (..),+    DistStm (..),+    DistStms,+    DistBody (..),+    DistInput (..),+    DistInputs,+    DistType (..),+    distInputType,+    DistResult (..),+    ResTag (..),+    DistIrregularity (..),+    FunHasParallelism,+    isRegularDistResult,+    isParallelStm,+    stmHasMeaningfulParallelism,++    -- * Segments+    Segments,+    segmentsShape,+    segmentsRank,+    segmentCount,+  )+where++import Data.Bifunctor+import Data.Foldable+import Data.List qualified as L+import Data.Map qualified as M+import Data.Maybe+import Data.Sequence qualified as Seq+import Data.Set qualified as S+import Futhark.Analysis.PrimExp.Convert+import Futhark.IR.SOACS+import Futhark.Util (nubOrd)+import Futhark.Util.Pretty++-- | Widths of the enclosing map-nest, outermost first. For top-level parallel+-- constructs, this is empty, which should be treated as an implicit+-- single-element segment (see 'segmentCount'). Generally, the empty-segments+-- case must be treated specially in some places, which is unfortunate, but it+-- helps unify code between the nested and top level cases.+type Segments = [SubExp]++type FunHasParallelism = Name -> Bool++-- | How to treat irregularity when classifying the statements of a distributed+-- body. This is mainly used to sequentialise nonuniform nested parallelism+-- instead of actually exploiting the parallelism, as the overhead of doing so+-- can sometimes be ruinous.+data DistIrregularity+  = -- | Distribute statements involving irregularity, relying on the machinery+    -- for flattening irregular arrays to handle them.+    DistributeIrregular+  | -- | Sequentialise BasicOps that involve only internal nonuniformity instead+    -- of distributing them. Used when generating intrablock code, where the+    -- machinery for flattening nonuniform nested parallelism would produce+    -- SegOps whose sizes are bound inside the kernel body, which makes the+    -- enclosing intrablock kernel infeasible ('noNonuniformPar' would reject+    -- it). Irregularity that escapes the enclosing map must still be+    -- distributed; if it occurs, the intrablock version is correctly rejected.+    SequentialiseIrregularBasicOps+  | -- | Sequentialise /any/ statement whose nonuniformity stays internal.+    SequentialiseIrregularAll+  deriving (Eq, Show)++segmentsShape :: Segments -> Shape+segmentsShape = Shape++segmentsRank :: Segments -> Int+segmentsRank = shapeRank . segmentsShape++segmentCount :: Segments -> TPrimExp Int64 VName+segmentCount = product . map pe64 . shapeDims . segmentsShape++newtype ResTag = ResTag Int+  deriving (Eq, Ord, Show)++-- | Something that is mapped.+data DistInput+  = -- | A value bound outside the original map-nest.  By necessity+    -- regular.  The type is the parameter type.+    DistInputFree VName Type+  | -- | A value constructed inside the original map-nest.  May be+    -- irregular.+    DistInput ResTag Type+  deriving (Eq, Ord, Show)++type DistInputs = [(VName, DistInput)]++nubInputs :: DistInputs -> DistInputs+nubInputs = L.nubBy (\a b -> fst a == fst b)++-- | The type of a 'DistInput'.  This corresponds to the parameter+-- type of the original map-nest.+distInputType :: DistInput -> Type+distInputType (DistInputFree _ t) = t+distInputType (DistInput _ t) = t++data DistType+  = DistType+      -- | Outer regular size.+      Segments+      -- | Irregular dimensions on top (but after the leading regular+      -- size).+      Rank+      -- | The regular "element type" - in the worst case, at least a+      -- scalar.+      Type+  deriving (Eq, Ord, Show)++data DistResult = DistResult {distResTag :: ResTag, distResType :: DistType, distResName :: VName}+  deriving (Eq, Ord, Show)++-- | The body of a distributed statement.+data DistBody+  = -- | A single statement that may involve parallel operations or produce an+    -- irregular array.+    ParallelStm (Stm SOACS)+  | -- | Single or Multiple scalar operations grouped into a single traversal+    ScalarStm (Stms SOACS)+  deriving (Eq, Ord, Show)++distBodyStms :: DistBody -> Stms SOACS+distBodyStms (ParallelStm stm) = oneStm stm+distBodyStms (ScalarStm stms) = stms++data DistStm = DistStm+  { distStmInputs :: DistInputs,+    distStmResult :: [DistResult],+    distStmBody :: DistBody+  }+  deriving (Eq, Ord, Show)++distStmStms :: DistStm -> Stms SOACS+distStmStms = distBodyStms . distStmBody++-- | An efficient sequence of 'DistStm's.+type DistStms = Seq.Seq DistStm++-- | First element of tuple are certificates for this result.+--+-- Second is the name to which is should be bound.+--+-- Third is the element type (i.e. excluding shape of segments).+type ResMap = M.Map ResTag [([DistInput], VName, Type)]++-- | The results of a map-distribution that were free or identity+-- mapped in the original map function.  These correspond to plain+-- replicated arrays.+type DistRep = (VName, Either SubExp DistInput)++data DistResults = DistResults ResMap [DistRep]+  deriving (Eq, Ord, Show)++data Distributed = Distributed DistStms DistResults+  deriving (Eq, Ord, Show)++instance Pretty ResTag where+  pretty (ResTag x) = "r" <> pretty x++instance Pretty DistInput where+  pretty (DistInputFree v _) = pretty v+  pretty (DistInput rt _) = pretty rt++instance Pretty DistType where+  pretty (DistType w r t) =+    brackets (pretty w) <> pretty r <> pretty t++instance Pretty DistResult where+  pretty (DistResult rt t _) =+    pretty rt <> colon <+> pretty t++instance Pretty DistStm where+  pretty (DistStm inputs res stms) =+    "let" <+> ppTuple' (map pretty res) <+> "=" </> indent 2 stm'+    where+      stm' =+        "map"+          <+> nestedBlock+            ( stack $+                map onInput inputs+                  ++ map pretty (toList (distBodyStms stms))+                  ++ [ "return" <+> ppTuple' (map pretty res)+                     ]+            )+      onInput (v, inp) =+        "for"+          <+> parens (pretty v <> colon <+> pretty (distInputType inp))+          <+> "<-"+          <+> pretty inp++instance Pretty Distributed where+  pretty (Distributed stms (DistResults resmap reps)) =+    stms' </> res'+    where+      res' = stack $ map onRes (M.toList resmap) <> map onRep reps+      stms' = stack $ map pretty $ toList stms+      onRes (rt, binds) =+        stack ["let" <+> pretty v <+> "=" <+> pretty rt | v <- binds]+      onRep (v, Left se) =+        "let" <+> pretty v <+> "=" <+> "rep" <> parens (pretty se)+      onRep (v, Right tag) =+        "let" <+> pretty v <+> "=" <+> "rep" <> parens (pretty tag)++resultMap :: [(VName, DistInput)] -> DistStms -> Pat Type -> Result -> ResMap+resultMap avail_inputs stms pat res = foldMap (foldMap f . distStmResult) stms+  where+    pes = M.fromList $ do+      stm <- toList stms+      pe <- concatMap (patElems . stmPat) (distStmStms stm)+      pure (patElemName pe, pe)+    f (DistResult rt _ v) =+      case maybe [] findRess $ M.lookup v pes of+        [] -> mempty+        binds -> M.singleton rt binds+    findRess (PatElem v v_t) = do+      (SubExpRes cs se, pv) <- zip res (patNames pat)+      if se == Var v+        then pure (map findCert (unCerts cs), pv, v_t)+        else []+    findCert v = fromMaybe (DistInputFree v (Prim Unit)) $ lookup v avail_inputs++splitIrregDims :: Names -> Type -> (Rank, Type)+splitIrregDims bound_outside (Array pt shape u) =+  let (reg, irreg) =+        first reverse $ span regDim $ reverse $ shapeDims shape+   in (Rank $ length irreg, Array pt (Shape reg) u)+  where+    regDim (Var v) = v `nameIn` bound_outside+    regDim Constant {} = True+splitIrregDims _ t = (mempty, t)++freeInput :: [(VName, DistInput)] -> VName -> Maybe (VName, DistInput)+freeInput avail_inputs v =+  (v,) <$> lookup v avail_inputs++patInput :: ResTag -> PatElem Type -> (VName, DistInput)+patInput tag pe =+  (patElemName pe, DistInput tag $ patElemType pe)++nextResTag :: DistInputs -> ResTag+nextResTag = foldl' step (ResTag 0)+  where+    step next (_, DistInputFree _ _) =+      next+    step next (_, DistInput (ResTag i) _) =+      max next (ResTag (i + 1))++distributeBody ::+  DistIrregularity ->+  FunHasParallelism ->+  Scope rep ->+  Segments ->+  DistInputs ->+  Body SOACS ->+  (DistInputs, DistStms)+distributeBody irreg_mode funHasParallelism outer_scope w param_inputs body = do+  let ((_, avail_inputs), stms) =+        L.mapAccumL distributeStm (nextResTag param_inputs, param_inputs) $+          bodyStms body+   in ( avail_inputs,+        classifyStms irreg_mode funHasParallelism (bodyResult body) stms+      )+  where+    bound_outside = namesFromList $ M.keys outer_scope+    distType t = uncurry (DistType w) $ splitIrregDims bound_outside t+    distributeStm (ResTag tag, avail_inputs) stm =+      let pat = stmPat stm+          new_tags = map ResTag $ take (patSize pat) [tag ..]+          avail_inputs' =+            avail_inputs <> zipWith patInput new_tags (patElems pat)+          free_in_stm = freeIn stm+          used_free = mapMaybe (freeInput avail_inputs) $ namesToList free_in_stm+          used_free_types =+            mapMaybe (freeInput avail_inputs)+              . namesToList+              . foldMap (freeIn . distInputType . snd)+              $ used_free+          stm' =+            DistStm+              (nubInputs $ used_free_types <> used_free)+              (zipWith3 DistResult new_tags (map distType $ patTypes pat) (patNames pat))+              (ParallelStm stm)+       in ((ResTag $ tag + length new_tags, avail_inputs'), stm')++isParallelDistStm :: (Stm SOACS -> Bool) -> DistStm -> Bool+isParallelDistStm stm_is_parallel (DistStm _ res (ParallelStm stm)) =+  stm_is_parallel stm || not (all isRegularDistResult res)+isParallelDistStm _ _ = False++noSequentialAttr :: Stm SOACS -> Bool+noSequentialAttr stm =+  not ("sequential" `inAttrs` stmAuxAttrs (stmAux stm))++-- | Does the statement contain meaningful parallelism - a SOAC or a call to+-- a parallel function, possibly nested inside sequential control flow?+-- Basic operations such as 'Iota' or 'Replicate' do not count. See Note+-- [Meaningful Parallelism].+stmHasMeaningfulParallelism :: FunHasParallelism -> Stm SOACS -> Bool+stmHasMeaningfulParallelism funHasParallelism = hasParallelism+  where+    hasParallelism stm' =+      noSequentialAttr stm'+        && case stmExp stm' of+          BasicOp _ -> False+          Apply fname _ _ _ -> funHasParallelism fname+          Match _ cases def_case _ ->+            any hasParallelism $+              bodyStms def_case+                <> mconcat (map (bodyStms . caseBody) cases)+          Loop _ _ body -> any hasParallelism (bodyStms body)+          WithAcc _ lam -> any hasParallelism (bodyStms (lambdaBody lam))+          Op op -> isParallelOp op++    isParallelOp JVP {} = error "stmHasMeaningfulParallelism: JVP"+    isParallelOp VJP {} = error "stmHasMeaningfulParallelism: VJP"+    isParallelOp _ = True++isParallelStm :: FunHasParallelism -> Stm SOACS -> Bool+isParallelStm funHasParallelism stm =+  noSequentialAttr stm+    && (parallelBasicOp (stmExp stm) || stmHasMeaningfulParallelism funHasParallelism stm)+  where+    parallelBasicOp (BasicOp op) = isParallelBasicOp op+    parallelBasicOp _ = False++    isParallelBasicOp (Update _ _ slice _) = not $ null $ sliceDims slice+    isParallelBasicOp Concat {} = True+    isParallelBasicOp Iota {} = True+    isParallelBasicOp Replicate {} = True+    isParallelBasicOp (FlatUpdate _ flat_slice _) = not $ null $ flatSliceDims flat_slice+    isParallelBasicOp Manifest {} = True+    isParallelBasicOp Rearrange {} = True+    isParallelBasicOp Reshape {} = True+    isParallelBasicOp (FlatIndex _ flat_slice) = not $ null $ flatSliceDims flat_slice+    isParallelBasicOp (Index _ slice) = not $ null $ sliceDims slice+    -- Now the sequential ones - we handle them explicitly so we will notice if+    -- we ever add a new one.+    isParallelBasicOp ArrayLit {} = False+    isParallelBasicOp ArrayVal {} = False+    isParallelBasicOp Scratch {} = False+    isParallelBasicOp SubExp {} = False+    isParallelBasicOp Opaque {} = False+    isParallelBasicOp UnOp {} = False+    isParallelBasicOp BinOp {} = False+    isParallelBasicOp CmpOp {} = False+    isParallelBasicOp ConvOp {} = False+    isParallelBasicOp Assert {} = False+    isParallelBasicOp UpdateAcc {} = False+    isParallelBasicOp UserParam {} = False++isRegularDistResult :: DistResult -> Bool+isRegularDistResult (DistResult _ (DistType _ (Rank r) _) _) = r == 0++-- | Does the statement contain, inside a 'Loop' or 'Match', a statement whose+-- result is sized by a name bound within the enclosing statement itself? Such a+-- statement can never be executed sequentially inside a GPU kernel, as it+-- implies an allocation whose size cannot be computed before the kernel is+-- launched; only the machinery for flattening irregular arrays can handle it.+-- In contrast, sizes that are nonuniform merely by being variant to the+-- enclosing map-nest are fine, as memory expansion can compute those via+-- slicing. This is essentially a heuristic where we bet that slicing is+-- efficient; the fully principled stance would be to not allow any+-- nonuniformity. This is one of the criteria of 'mustDistribute' in+-- 'classifyStms'. See Note [Meaningful Parallelism].+stmHasNonuniformInside :: Stm SOACS -> Bool+stmHasNonuniformInside = inExp mempty . stmExp+  where+    nonuniform bound = any (`nameIn` bound) . subExpVars . arrayDims+    inExp bound e =+      case e of+        Loop merge form body ->+          inBody+            ( bound+                <> namesFromList (map (paramName . fst) merge)+                <> namesFromList (M.keys (scopeOfLoopForm form))+            )+            body+        Match _ cases def_body _ ->+          any (inBody bound) (def_body : map caseBody cases)+        WithAcc _ lam -> inBody bound (lambdaBody lam)+        _ -> False+    inBody bound body =+      any (inStm (bound <> boundInBody body)) $ bodyStms body+    inStm bound stm =+      any (nonuniform bound) (patTypes (stmPat stm))+        || inExp bound (stmExp stm)++-- | Merge a group of scalar 'DistStm's into a single one.+mergeGroup :: Result -> DistStms -> DistStms -> DistStm+mergeGroup bodyRes ds rest =+  let resTags =+        S.fromList $ concatMap (map distResTag . distStmResult) ds+      isInternal (_, DistInput rt _) = rt `S.member` resTags+      isInternal _ = False+      externalInputs =+        nubInputs $+          concatMap (filter (not . isInternal) . distStmInputs) ds+      externalResults =+        nubOrd $+          concatMap (filter (isExternal bodyRes rest) . distStmResult) ds+      allStms = foldMap distStmStms ds+   in DistStm externalInputs externalResults (ScalarStm allStms)++groupStms ::+  (DistStm -> Bool) ->+  Result ->+  Seq.Seq DistStm ->+  Seq.Seq DistStm+groupStms _ _ Seq.Empty = mempty+groupStms dist_stm_is_parallel body_res ds' =+  let (scalars, rest) = Seq.breakl dist_stm_is_parallel ds'+      scalar_grouped+        | not $ null scalars =+            Seq.singleton $ mergeGroup body_res scalars rest+        | otherwise = mempty+   in case rest of+        Seq.Empty -> scalar_grouped+        p Seq.:<| ps ->+          scalar_grouped <> (p Seq.<| groupStms dist_stm_is_parallel body_res ps)++--  we should probably sort the DistStms first and we should assume they are sorted+-- and then given to this function.+classifyStms :: DistIrregularity -> FunHasParallelism -> Result -> DistStms -> DistStms+classifyStms irreg_mode funHasParallelism body_res = classify+  where+    -- Distribute the statements that are parallel, plus those that+    -- 'mustDistribute' regardless of parallelism. If no statement contains+    -- meaningful parallelism, no statement counts as parallel, so that the+    -- trivially parallel statements are treated as sequential as well. See Note+    -- [Meaningful Parallelism].+    --+    -- With 'SequentialiseIrregularBasicOps', nonuniform basic operations+    -- further count as parallel only when their irregular arrays escape the+    -- scalar group. With 'SequentialiseIrregularAll', this goes for any+    -- statement, not just basic operations.+    classify ds =+      let -- Which statements are candidates for sequentialising when their+          -- nonuniformity does not escape the scalar group.+          sequentialisable = case irreg_mode of+            DistributeIrregular -> const False+            SequentialiseIrregularBasicOps -> isBasicOpDistStm+            SequentialiseIrregularAll -> const True+          parallel+            | any meaningfulDistStm ds =+                \d ->+                  isParallelDistStm (isParallelStm funHasParallelism) d+                    && not (sequentialisable d && involvesIrregularity ds d)+            | otherwise = const False+          forced = mustDistribute (S.fromList $ filter parallel $ toList ds) ds+       in groupStms (\d -> parallel d || d `S.member` forced) body_res ds++    meaningfulDistStm d@(DistStm _ _ (ParallelStm stm)) =+      (isBasicOpDistStm d && any (`nameIn` freeIn body_res) (patNames (stmPat stm)))+        || stmHasMeaningfulParallelism funHasParallelism stm+    meaningfulDistStm _ = False++    isBasicOpDistStm (DistStm _ _ (ParallelStm (Let _ _ BasicOp {}))) = True+    isBasicOpDistStm _ = False++    distStmHasNonuniformInside (DistStm _ _ (ParallelStm stm)) =+      stmHasNonuniformInside stm+    distStmHasNonuniformInside _ = False++    -- Whether a statement produces an irregular array, or consumes+    -- one produced elsewhere in the body.+    involvesIrregularity ds d =+      not (all isRegularDistResult (distStmResult d))+        || any consumesIrregular (distStmInputs d)+      where+        consumesIrregular (_, DistInput rt _) = rt `S.member` irregular_tags+        consumesIrregular _ = False+        irregular_tags =+          S.fromList $+            map distResTag $+              filter (not . isRegularDistResult) $+                foldMap distStmResult ds++    -- The statements that require distribution regardless of whether they+    -- contain profitable parallelism, because a sequentially executed scalar+    -- group cannot handle their sizes:+    --+    -- (1) Statements whose irregular results are used outside the scalar+    --     group (by the body result, or transitively by another distributed+    --     statement), as arrays produced by sequentially executed groups+    --     must be regular.+    --+    -- (2) Compound statements (e.g. loops) with irregular results, even+    --     internally used ones: they contain allocations of nonuniform size+    --     that only flattening can handle, while a basic operation can+    --     reasonably be executed sequentially by a single thread when its+    --     result stays internal. This criterion does not apply under+    --     'SequentialiseIrregularAll', which is only imposed on user request,+    --     and may result in non-compileable code due to impossible memory+    --     expansion.+    --+    -- (3) Statements with sizes bound inside their own sequential control+    --     flow ('stmHasNonuniformInside').+    --+    -- The initial set of forced statements may be seeded with the statements+    -- already known to be distributed, such that irregular arrays they consume+    -- are also distributed.+    mustDistribute seed ds = fixpoint seed+      where+        body_res_names = freeIn body_res+        irregularResults stm =+          filter (not . isRegularDistResult) $ distStmResult stm+        forcedTags forced =+          S.fromList+            [rt | stm <- S.toList forced, (_, DistInput rt _) <- distStmInputs stm]+        compoundIrregular stm =+          irreg_mode /= SequentialiseIrregularAll+            && not (isBasicOpDistStm stm)+            && not (null (irregularResults stm))+        isForced forced stm =+          compoundIrregular stm+            || distStmHasNonuniformInside stm+            || any+              ( \r ->+                  distResName r `nameIn` body_res_names+                    || distResTag r `S.member` forcedTags forced+              )+              (irregularResults stm)+        fixpoint forced =+          let forced' = seed <> S.fromList (filter (isForced forced) $ toList ds)+           in if forced' == forced then forced else fixpoint forced'++-- | A result is external if it is used by a subsequent 'DistStm' or by the body+-- result.+isExternal :: Result -> DistStms -> DistResult -> Bool+isExternal bodyRes rest (DistResult rt _ rn) =+  rt `S.member` usedByRest || rn `S.member` bodyResVars || rn `S.member` bodyResCerts+  where+    usedByRest =+      S.fromList+        [rt' | (_, DistInput rt' _) <- concatMap distStmInputs rest]+    bodyResVars =+      S.fromList $+        mapMaybe+          ( \(SubExpRes _ se) -> case se of+              Var v -> Just v+              _ -> Nothing+          )+          bodyRes+    bodyResCerts =+      S.fromList $+        concatMap (\(SubExpRes cs _) -> unCerts cs) bodyRes++-- | The input we are mapping over in 'distributeMap'.+data MapArray t+  = -- | A straightforward array passed in to a top-level map.+    MapArray VName Type+  | -- | Something more exotic - distribution will assign it a 'ResTag', but not+    -- do anything else. This is used to distributed nested maps whose inputs+    -- are produced in the outer nests.+    MapOther t Type++mapArrayRowType :: MapArray t -> Type+mapArrayRowType (MapArray _ t) = t+mapArrayRowType (MapOther _ t) = t++-- This is used to handle those results that are constants or lambda+-- parameters.+findReps :: [(VName, DistInput)] -> Pat Type -> Lambda SOACS -> [DistRep]+findReps avail_inputs map_pat lam =+  mapMaybe f $ zip (patElems map_pat) (bodyResult (lambdaBody lam))+  where+    f (pe, SubExpRes _ (Var v)) =+      case lookup v avail_inputs of+        Nothing -> Just (patElemName pe, Left $ Var v)+        Just inp+          | v `elem` map paramName (lambdaParams lam) ->+              Just (patElemName pe, Right inp)+          | otherwise -> Nothing+    f (pe, SubExpRes _ (Constant v)) = do+      Just (patElemName pe, Left $ Constant v)++distributeMap ::+  DistIrregularity ->+  FunHasParallelism ->+  Scope rep ->+  Pat Type ->+  Segments ->+  [MapArray t] ->+  Lambda SOACS ->+  (Distributed, M.Map ResTag t)+distributeMap irreg_mode funHasParallelism outer_scope map_pat w arrs lam =+  let ((_, arrmap), param_inputs) =+        L.mapAccumL paramInput (ResTag 0, mempty) $+          zip (lambdaParams lam) arrs+      (avail_inputs, stms) =+        distributeBody irreg_mode funHasParallelism outer_scope w param_inputs $ lambdaBody lam+      resmap =+        resultMap avail_inputs stms map_pat $+          bodyResult (lambdaBody lam)+      reps = findReps avail_inputs map_pat lam+   in ( Distributed stms $ DistResults resmap reps,+        arrmap+      )+  where+    paramInput (ResTag i, m) (p, MapArray arr _) =+      ( (ResTag i, m),+        (paramName p, DistInputFree arr $ paramType p)+      )+    paramInput (ResTag i, m) (p, MapOther x _) =+      ( (ResTag (i + 1), M.insert (ResTag i) x m),+        (paramName p, DistInput (ResTag i) $ paramType p)+      )++-- Note [Meaningful Parallelism]+--+-- Many basic operations (Iota, Replicate, Concat, nontrivial slicing, and so+-- on) are parallel in principle, but distributing one on its own gains us+-- nothing: it performs the same work as a sequential per-thread traversal, at+-- the cost of manifesting intermediate arrays and launching extra kernels. Our+-- rule of thumb is that code contains *meaningful* parallelism only when it+-- contains at least a Screma, Hist, or a call to a parallel function - possibly+-- nested inside sequential control flow. This is what+-- 'stmHasMeaningfulParallelism' checks. It is intended to avoid "parallelising"+-- map bodies that are essentially sequential loops, but happen to use some+-- basic operations to construct arrays they operate on. This is a somewhat+-- crude classification, and it is possible to imagine a more sophisticated one+-- based on a cost-model. Note that incremental flattening (and auto-tuning)+-- might allow us to generate and pick the sequentialised versions anyway, but+-- it is more efficient to not generate them in the first place.+--+-- The notion is used in two places:+--+-- (1) When classifying the statements of a distributed body as parallel or+--     sequential ('classifyStms'): if no statement contains meaningful+--     parallelism, then the trivially parallel statements are also grouped+--     as sequential, so that ideally the entire body becomes a single+--     segmented operation.+--+-- (2) When deciding whether to multi-version a Screma+--     ('factorScremaForParallelism' in Futhark.Pass.Flatten.Incremental): a+--     fully flattened code version is worth generating only when the lambda+--     contains meaningful parallelism.+--+-- One exception is nonuniformity. A statement whose result shape varies+-- across the surrounding map-nest (e.g. 'iota x' for a mapped 'x') cannot+-- in general be traversed sequentially per thread, as arrays produced as+-- results of sequentially executed groups must be regular - handling it+-- is exactly what flattening is for. Hence irregular results still count+-- as meaningful, with one refinement in (1): if the irregular arrays are+-- used only *within* the scalar group (e.g. slices feeding a concatenation+-- whose size is ultimately uniform), the statements can be executed+-- sequentially after all - see 'mustDistribute' in 'classifyStms'. In (2)+-- irregular results make versioning worthwhile+-- ('lambdaHasMeaningfulParallelism' in Futhark.Pass.Flatten.Incremental).+--+-- The refinement above does not extend to nonuniformity arising *inside* a+-- sequential Loop or Match: an array sized by a loop-variant name implies an+-- allocation whose size cannot be computed before a kernel is launched, so a+-- statement containing one can never be part of a sequentially executed+-- scalar group, even if nothing irregular escapes it. Such statements are+-- unconditionally forced by 'mustDistribute' - see 'stmHasNonuniformInside'.+--+-- Another exception is when a trivial statement (a basic operation) produces+-- the result of a body - in that case we also treat it as meaningful, because+-- we have to manifest it. This does not extend to compound statements (e.g. a+-- sequential loop) producing a body result: those contain no parallelism worth+-- distributing on their own, and treating them as meaningful would needlessly+-- split any surrounding sequential statements into multiple kernels. This is+-- the part that would benefit from a cost model, or simply from being more+-- principled and ignoring the small inefficiencies.
+ src/Futhark/Pass/Flatten/General.hs view
@@ -0,0 +1,884 @@+-- | Various general utilities used in flattening.+module Futhark.Pass.Flatten.General+  ( -- * Reading inputs+    readInputVar,+    readInputs,+    readInput,+    readNeutral,+    readTypeDims,++    -- * Building blocks+    ensureDenseIrregular,+    liftResult,+    liftDistResultRep,+    liftSubExp,+    liftSubExpPreserveRep,+    liftSubExpRegular,+    liftVarRegular,+    isRegularInputArr,+    liftParam,+    liftRegularParam,+    liftRegResult,+    needsIrregularRetType,+    mkIrregFromReg,+    flattenIrregularRep,+    distCerts,+    dataArr,+    getIrregRep,+    scatterIrregular,+    scatterRegular,+    module Futhark.Pass.Flatten.Monad,+    module Futhark.Pass.Flatten.Builtins,++    -- * Various+    scopeOfDistInputs,+    lookupInputType,+    subExpInputType,+    localiseInputs,+    replicateForDims,+    liftBodyWithDistResults,+    distResultsToResReps,+    resultToResReps,+    resultToResRepsByDistResult,+    irregularRepToFlatArrs,+    distributeAndFlattenBody,+    splitInput,+    isVariant,+    segmentDims,+    flattenDistStm,+    flattenScalarStm,+    distributeBodyWith,+    distributeMapWith,+    atSegLevel,+  )+where++import Control.Monad+import Data.List qualified as L+import Data.Map qualified as M+import Data.Maybe+import Data.Tuple.Solo+import Futhark.IR.GPU+import Futhark.IR.SOACS as SOACS+import Futhark.MonadFreshNames+import Futhark.Pass.Flatten.Builtins+import Futhark.Pass.Flatten.Distribute+import Futhark.Pass.Flatten.Monad+import Futhark.Tools+import Futhark.Transform.Rename (renameExp)+import Futhark.Util (mapAccumLM)+import Futhark.Util.IntegralExp+import Prelude hiding (div, rem)++-- | Write back the irregular results of a branch to a (partially) blank space+-- The `offsets` variable is the offsets of the final result, whereas `irregRep`+-- is the irregular representation of the result.+scatterIrregular ::+  SegLevel ->+  VName ->+  VName ->+  (VName, IrregularRep) ->+  FlattenM VName+scatterIrregular lvl offsets space (is, irregRep) = do+  let IrregularRep {irregularS = segs, irregularD = elems, irregularO = off, irregularK = kind} = irregRep+  (_, _, ii1) <- doRepIota lvl segs+  (_, _, ii2) <- doSegIota lvl segs+  m <- arraySize 0 <$> lookupType ii1+  letExp "irregular_scatter" <=< genScatter lvl space m $ \gtid -> do+    segment <- letSubExp "segment" =<< eIndex ii1 [eSubExp gtid]+    intra_segment <- letSubExp "segment" =<< eIndex ii2 [eSubExp gtid]+    x <- case kind of+      Dense -> letSubExp "x" =<< eIndex elems [eSubExp gtid]+      Replicated -> do+        o <- letSubExp "rep_O" =<< eIndex off [eSubExp segment]+        letSubExp "x" =<< eIndex elems [toExp $ pe64 o + pe64 intra_segment]+    offset <- letExp "offset" =<< eIndex offsets [eIndex is [eSubExp segment]]+    i <- letExp "i" =<< eBinOp (Add Int64 OverflowUndef) (toExp offset) (eSubExp intra_segment)+    pure (i, x)++-- | Write back the regular results to a (partially) blank space+scatterRegular ::+  SegLevel ->+  VName ->+  (VName, VName) ->+  FlattenM VName+scatterRegular lvl space (is, xs) = do+  dims <- arrayDims <$> lookupType xs+  letExp "regular_scatter" <=< genScatterND lvl space dims $ \(gtid : rest) -> do+    x <- letSubExp "x" =<< eIndex xs (map eSubExp (gtid : rest))+    i <- letSubExp "i" =<< eIndex is [eSubExp gtid]+    pure (i : rest, x)++ensureDenseIrregular :: SegLevel -> Name -> IrregularRep -> FlattenM IrregularRep+ensureDenseIrregular _ _ rep@IrregularRep {irregularK = Dense} =+  pure rep+ensureDenseIrregular lvl desc rep@IrregularRep {} = do+  (new_F, new_O, ii1) <- doRepIota lvl (irregularS rep)+  m <- arraySize 0 <$> lookupType ii1+  new_D <- letExp (desc <> "_dense_D") <=< segMap lvl (MkSolo m) $ \(MkSolo i) -> do+    seg <- letSubExp "seg" =<< eIndex ii1 [eSubExp i]+    old_off <- letSubExp "old_off" =<< eIndex (irregularO rep) [eSubExp seg]+    new_off <- letSubExp "new_off" =<< eIndex new_O [eSubExp seg]+    j <- letSubExp "j" <=< toExp $ pe64 i - pe64 new_off+    x <- letSubExp "x" =<< eIndex (irregularD rep) [toExp $ pe64 old_off + pe64 j]+    pure [subExpRes x]+  pure $+    IrregularRep+      { irregularS = irregularS rep,+        irregularF = new_F,+        irregularO = new_O,+        irregularD = new_D,+        irregularK = Dense+      }++-- Lift a result of a function.+liftResult :: SegLevel -> Segments -> DistInputs -> DistEnv -> SubExpRes -> FlattenM Result+liftResult lvl segments inps env res = map (SubExpRes mempty . Var) <$> vs+  where+    vs = do+      (_, rep) <- liftSubExp lvl segments inps env (resSubExp res)+      case rep of+        Regular v -> pure [v]+        Irregular irreg -> mkIrrep irreg+    mkIrrep+      ( IrregularRep+          { irregularS = segs,+            irregularF = flags,+            irregularO = offsets,+            irregularD = elems+          }+        ) = do+        flags_t <- lookupType flags+        t <- lookupType elems+        num_data <- letExp "num_data" =<< toExp (product $ map pe64 $ arrayDims t)+        let shape = Shape [Var num_data]+        flags' <- letExp "flags" $ BasicOp $ Reshape flags $ reshapeAll (arrayShape flags_t) shape+        elems' <- letExp "elems" $ BasicOp $ Reshape elems $ reshapeAll (arrayShape t) shape+        pure [num_data, segs, flags', offsets, elems']++needsIrregularRetType :: DistInputs -> RetType SOACS -> Bool+needsIrregularRetType inps = any needsIrregularDim . arrayDims+  where+    needsIrregularDim Ext {} = True+    needsIrregularDim (Free se) = isVariant inps se++liftRegResult :: SegLevel -> Segments -> SubExp -> DistInputs -> DistEnv -> RetType SOACS -> SubExpRes -> FlattenM Result+liftRegResult lvl segments num_segments inps env rettype res+  | needsIrregularRetType inps rettype = case resSubExp res of+      Var v -> do+        irreg <- getIrregRep lvl segments env inps v+        varsRes <$> irregularRepToFlatArrs num_segments irreg+      Constant {} ->+        error "liftRegResult: irregular result is not a variable"+  | otherwise = do+      let res_se = resSubExp res+      res_t <- subExpInputType inps res_se+      let expectedShape = segmentsShape segments <> arrayShape res_t+      lifted_res <- liftSubExpRegular lvl segments inps env expectedShape res_se+      pure [SubExpRes mempty (Var lifted_res)]++mkIrregFromReg ::+  SegLevel ->+  Segments ->+  VName ->+  FlattenM IrregularRep+mkIrregFromReg lvl segments arr = do+  arr_t <- lookupType arr+  num_segments <-+    letSubExp "reg_num_segments" <=< toExp $ product $ segmentDims segments+  segment_size <-+    letSubExp "reg_seg_size" <=< toExp . product . map pe64 $+      drop (segmentsRank segments) (arrayDims arr_t)+  arr_S <-+    letExp "reg_segments" . BasicOp $+      Replicate (Shape [num_segments]) segment_size+  num_elems <-+    letSubExp "reg_num_elems" <=< toExp $ product $ map pe64 $ arrayDims arr_t+  arr_D <-+    letExp "reg_D" . BasicOp $+      Reshape arr (reshapeAll (arrayShape arr_t) (Shape [num_elems]))+  arr_F <- letExp "reg_F" <=< segMap lvl (MkSolo num_elems) $ \(MkSolo i) -> do+    flag <- letSubExp "flag" <=< toExp $ (pe64 i `rem` pe64 segment_size) .==. 0+    pure [subExpRes flag]+  arr_O <- letExp "reg_O" <=< segMap lvl (MkSolo num_segments) $ \(MkSolo i) -> do+    offset <- letSubExp "offset" <=< toExp $ pe64 i * pe64 segment_size+    pure [subExpRes offset]+  pure $+    IrregularRep+      { irregularS = arr_S,+        irregularF = arr_F,+        irregularO = arr_O,+        irregularD = arr_D,+        irregularK = Dense+      }++readIrregularInput ::+  Segments ->+  [SubExp] ->+  VName ->+  Type ->+  IrregularRep ->+  FlattenM VName+readIrregularInput segments is v t (IrregularRep _ _ v_O v_D _) = do+  offset <- letSubExp "offset" =<< eIndex v_O [toExp $ flatSegmentIndex segments is]+  case arrayDims t of+    [] -> do+      letExp (baseName v <> "_inp") =<< eIndex v_D [eSubExp offset]+    [num_elems] -> do+      let slice = Slice [DimSlice offset num_elems (intConst Int64 1)]+      letExp (baseName v <> "_inp") $ BasicOp $ Index v_D slice+    _ -> do+      num_elems <-+        letSubExp "num_elems" =<< toExp (product $ map pe64 $ arrayDims t)+      let slice = Slice [DimSlice offset num_elems (intConst Int64 1)]+      v_flat <-+        letExp (baseName v <> "_flat") $ BasicOp $ Index v_D slice+      v_flat_t <- lookupType v_flat+      letExp (baseName v <> "_inp") . BasicOp $+        Reshape v_flat (reshapeAll (arrayShape v_flat_t) (arrayShape t))++readInputVar :: Segments -> DistEnv -> [SubExp] -> DistInputs -> VName -> FlattenM VName+readInputVar segments env is inputs v =+  case lookup v inputs of+    Nothing -> pure v+    Just (DistInputFree arr t)+      | isAcc t -> pure arr+      | otherwise -> letExp (baseName v) =<< eIndex arr (map eSubExp is)+    Just (DistInput rt t) -> do+      case resVar rt env of+        Regular arr+          | isAcc t -> pure arr+          | otherwise -> letExp (baseName v) =<< eIndex arr (map eSubExp is)+        Irregular irreg -> readIrregularInput segments is v t irreg++readInput :: Segments -> DistEnv -> [SubExp] -> DistInputs -> SubExp -> FlattenM SubExp+readInput _ _ _ _ (Constant x) =+  pure $ Constant x+readInput segments env is inputs (Var v) =+  Var <$> readInputVar segments env is inputs v++-- | Read the neutral element of a reduction or scan operator. The neutral+-- element of an operator is unique because we assume uniform operators, so a+-- valid program must have the same value in every segment, and we can read it+-- from the first one. This is the case even if the neutral element is+-- nonuniform, which is useless but allowed. The segment space may however be+-- empty, in which case we produce a blank value instead; it is never used then,+-- but an unconditional read would be out of bounds.+readNeutral :: Segments -> DistEnv -> DistInputs -> SubExp -> FlattenM SubExp+readNeutral segments env inps ne+  | Var v <- ne,+    Just _ <- lookup v inps = do+      ne_t <- subExpInputType inps ne+      n <- letSubExp "num_segments" =<< toExp (segmentCount segments)+      letSubExp (baseName v <> "_ne")+        =<< eIf+          (toExp $ pe64 n .==. 0)+          (eBody [eBlank ne_t])+          (eBody [eSubExp =<< readInput segments env zeros inps ne])+  | otherwise =+      readInput segments env zeros inps ne+  where+    zeros = replicate (segmentsRank segments) (intConst Int64 0)++readTypeDims ::+  Segments ->+  DistEnv ->+  [SubExp] ->+  DistInputs ->+  TypeBase Shape u ->+  FlattenM [SubExp]+readTypeDims segments env is inputs =+  mapM (readInput segments env is inputs) . arrayDims++segmentDims :: Segments -> [TPrimExp Int64 VName]+segmentDims = map pe64 . shapeDims . segmentsShape++flatSegmentIndex :: Segments -> [SubExp] -> TPrimExp Int64 VName+flatSegmentIndex segments = flattenIndex (segmentDims segments) . map pe64++readInputs :: Segments -> DistEnv -> [SubExp] -> DistInputs -> FlattenM ()+readInputs segments env is = mapM_ onInput+  where+    bindInputName v e+      | v `nameIn` freeIn e = do+          v' <- letExp (baseName v <> "_inp") e+          letBindNames [v] $ BasicOp $ SubExp $ Var v'+      | otherwise =+          letBindNames [v] e+    onInput (v, DistInputFree arr t) =+      bindInputName v+        =<< if isAcc t+          then eSubExp (Var arr)+          else eIndex arr (map eSubExp is)+    onInput (v, DistInput rt t) =+      case resVar rt env of+        Regular arr ->+          bindInputName v+            =<< if isAcc t+              then eSubExp $ Var arr+              else eIndex arr (map eSubExp is)+        Irregular irreg ->+          readIrregularInput segments is v t irreg >>= eSubExp . Var >>= bindInputName v++scopeOfDistInputs :: DistInputs -> Scope GPU+scopeOfDistInputs = scopeOfLParams . map f+  where+    f (v, inp) = Param mempty v (distInputType inp)++lookupInputType :: DistInputs -> VName -> FlattenM Type+lookupInputType inps v =+  case lookup v inps of+    Just inp -> pure $ distInputType inp+    Nothing -> lookupType v++subExpInputType :: DistInputs -> SubExp -> FlattenM Type+subExpInputType _ (Constant val) =+  pure $ Prim $ primValueType val+subExpInputType inps (Var v) =+  lookupInputType inps v++isVariant :: DistInputs -> SubExp -> Bool+isVariant _ (Constant _) = False+isVariant inps (Var v) = isJust $ lookup v inps++liftDistResultRep ::+  SegLevel ->+  Segments ->+  DistInputs ->+  DistEnv ->+  DistResult ->+  SubExpRes ->+  FlattenM ResRep+liftDistResultRep lvl segments inps env dist_res res+  | isRegularDistResult dist_res = do+      let (DistType _ _ t) = distResType dist_res+          expectedShape = segmentsShape segments <> arrayShape t+      Regular <$> liftSubExpRegular lvl segments inps env expectedShape (resSubExp res)+  | otherwise =+      case resSubExp res of+        Var v -> do+          rep <- getIrregRep lvl segments env inps v+          pure $ Irregular rep+        _ -> error "liftBranchResultRep: irregular result is not a variable"++liftDistResult :: SegLevel -> Segments -> DistInputs -> DistEnv -> DistResult -> SubExpRes -> FlattenM Result+liftDistResult lvl segments inps env dist_res res =+  if isRegularDistResult dist_res+    then do+      let (DistType _ _ t) = distResType dist_res+      let expectedShape = segmentsShape segments <> arrayShape t+      v <- liftSubExpRegular lvl segments inps env expectedShape (resSubExp res)+      pure [SubExpRes mempty (Var v)]+    else case resSubExp res of+      Var v -> do+        irreg <- getIrregRep lvl segments env inps v+        pure $ map (SubExpRes mempty . Var) [irregularS irreg, irregularF irreg, irregularO irreg, irregularD irreg]+      _ -> error "liftDistResult: irregular result is not a variable"++liftBodyWithDistResults :: FlattenOps -> Segments -> DistInputs -> DistEnv -> DistStms -> [DistResult] -> Result -> FlattenM Result+liftBodyWithDistResults ops segments inputs env dstms dist_res result = do+  env' <- foldM (flattenDistStm ops segments) env dstms+  result' <- zipWithM (liftDistResult (flattenSegLevel ops) segments inputs env') dist_res result+  pure $ concat result'++distResultsToResReps :: [DistResult] -> [VName] -> [ResRep]+distResultsToResReps dist_res results =+  snd $ L.mapAccumL f results dist_res+  where+    f rs dist_res' =+      if isRegularDistResult dist_res'+        then+          let (v : rs') = rs+           in (rs', Regular v)+        else+          let (segs : flags : offsets : elems : rs') = rs+           in (rs', Irregular $ IrregularRep segs flags offsets elems Dense)++-- | Convert an irregular representation to its flat constituents (number of+-- data elements, segments, flags, offsets, elements), with the structure arrays+-- reshaped to be one-dimensional of the given width. This is the form in which+-- irregular values are passed to lifted functions and carried through loops.+irregularRepToFlatArrs :: SubExp -> IrregularRep -> FlattenM [VName]+irregularRepToFlatArrs w (IrregularRep segs flags offsets elems _) = do+  t <- lookupType elems+  t_o <- lookupType offsets+  flags_t <- lookupType flags+  num_data <- letExp "num_data" =<< toExp (product $ map pe64 $ arrayDims t)+  let shape = Shape [Var num_data]+  flags' <- letExp "flags" $ BasicOp $ Reshape flags $ reshapeAll (arrayShape flags_t) shape+  elems' <- letExp "elems" $ BasicOp $ Reshape elems $ reshapeAll (arrayShape t) shape+  segs' <- letExp "segs" $ BasicOp $ Reshape segs $ reshapeAll (arrayShape t_o) (Shape [w])+  offsets' <- letExp "offsets" $ BasicOp $ Reshape offsets $ reshapeAll (arrayShape t_o) (Shape [w])+  pure [num_data, segs', flags', offsets', elems']++-- | Distribute a body and lift the distributed statements, giving+-- back representations of the body results.+distributeAndFlattenBody ::+  FlattenOps ->+  Segments ->+  Name ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  Body SOACS ->+  FlattenM [ResRep]+distributeAndFlattenBody ops segments desc env inps res body = do+  scope <- askScope+  (inps_local, env_local, _) <- localiseInputs env inps+  let (inps_dist, dstms) = distributeBodyWith ops scope segments inps_local body+  lifted_res <- liftBodyWithDistResults ops segments inps_dist env_local dstms res (bodyResult body)+  lifted_vs <- mapM (letExp desc <=< toExp . resSubExp) lifted_res+  pure $ distResultsToResReps res lifted_vs++-- | Take the elements at index @is@ from an input @v@. The representation of+-- @v@ can be overridden through the provided mapping, which takes precedence+-- over what the environment says.+splitInput ::+  SegLevel ->+  Segments ->+  DistEnv ->+  DistInputs ->+  VName ->+  M.Map VName ResRep ->+  VName ->+  FlattenM (Type, VName, ResRep)+splitInput lvl segments env inps is acc_reps v = do+  (t, rep0) <- liftSubExpPreserveRep segments inps env (Var v)+  let rep = M.findWithDefault rep0 v acc_reps+  (t,v,) <$> case rep of+    Regular arr -> do+      if isAcc t+        then+          pure $ Regular arr+        else do+          -- In the regular case we just take the elements+          -- of the array given by `is`+          n <- letSubExp "n" =<< (toExp . arraySize 0 =<< lookupType is)+          inner_dims <- drop (segmentsRank segments) . arrayDims <$> lookupType arr+          -- Do the segMap over all dims, so the inner dimensions+          -- are gathered in parallel+          arr' <- letExp "split_arr" <=< segMap lvl (n : inner_dims) $ \(i : js) -> do+            idx <- letSubExp "idx" =<< eIndex is [eSubExp i]+            let arr_is = unflattenIndex (segmentDims segments) (pe64 idx)+            subExpsRes . pure <$> (letSubExp "arr" =<< eIndex arr (map toExp arr_is ++ map eSubExp js))+          pure $ Regular arr'+    Irregular (IrregularRep segs flags offsets elems _) -> do+      -- In the irregular case we take the elements+      -- of the `segs` array given by `is` like in the regular case+      n <- letSubExp "n" =<< (toExp . arraySize 0 =<< lookupType is)+      segs' <- letExp "split_segs" <=< segMap lvl (MkSolo n) $ \(MkSolo i) -> do+        idx <- letExp "idx" =<< eIndex is [eSubExp i]+        subExpsRes . pure <$> (letSubExp "segs" =<< eIndex segs [toExp idx])+      -- From this we calculate the offsets and number of elements+      (_, offsets', num_data) <- exScanAndSum lvl segs'+      (_, _, ii1) <- doRepIota lvl segs'+      (_, _, ii2) <- doSegIota lvl segs'+      -- We then take the elements we need from `elems` and `flags`+      -- For each index `i`, we roughly:+      -- Get the offset of the segment we want to copy by indexing+      -- `offsets` through `is` further through `ii1` i.e.+      -- `offset = offsets[is[ii1[i]]]`+      -- We then add `ii2[i]` to `offset`+      -- and use that to index into `elems` and `flags`.+      ~[flags', elems'] <- letTupExp "split_F_data" <=< segMap lvl (MkSolo num_data) $ \(MkSolo i) -> do+        offset <- letExp "offset" =<< eIndex offsets [eIndex is [eIndex ii1 [eSubExp i]]]+        idx <- letExp "idx" =<< eBinOp (Add Int64 OverflowUndef) (toExp offset) (eIndex ii2 [eSubExp i])+        flags_split <- letSubExp "flags" =<< eIndex flags [toExp idx]+        elems_split <- letSubExp "elems" =<< eIndex elems [toExp idx]+        pure $ subExpsRes [flags_split, elems_split]+      pure $+        Irregular $+          IrregularRep+            { irregularS = segs',+              irregularF = flags',+              irregularO = offsets',+              irregularD = elems',+              irregularK = Dense+            }++-- | Flatten the arrays of an IrregularRep to be entirely one-dimensional.+flattenIrregularRep :: SegLevel -> IrregularRep -> FlattenM IrregularRep+flattenIrregularRep lvl ir@(IrregularRep shape _ offsets elems kind) = do+  elems_t <- lookupType elems+  if arrayRank elems_t == 1+    then pure ir+    else do+      n <- arraySize 0 <$> lookupType shape+      m' <- letSubExp "flat_m" <=< toExp $ product $ map pe64 $ arrayDims elems_t+      elems' <-+        letExp (baseName elems <> "_flat") . BasicOp $+          Reshape elems (reshapeAll (arrayShape elems_t) (Shape [m']))+      let inner_size = product $ map pe64 $ tail $ arrayDims elems_t+      ~[shape', offsets'] <-+        letTupExp (baseName shape <> "_flat_metadata")+          <=< renameExp+          <=< segMap lvl (MkSolo n)+          $ \(MkSolo i) -> do+            old_shape <-+              letSubExp "old_shape" =<< eIndex shape [toExp i]+            old_offset <-+              letSubExp "old_offset" =<< eIndex offsets [toExp i]++            segment_shape <-+              letSubExp "segment_shape" <=< toExp $ (pe64 old_shape * inner_size)+            segment_offset <-+              letSubExp "segment_offset" <=< toExp $ (pe64 old_offset * inner_size)+            pure $ subExpsRes [segment_shape, segment_offset]+      flags' <- genFlags lvl m' offsets'+      pure $ IrregularRep shape' flags' offsets' elems' kind++-- If the sub-expression is a constant, replicate it to match the shape of `segments`+-- If it's a variable, lookup the variable in the dist inputs and dist env,+-- and if it can't be found it is a free variable, so we replicate it to match the shape of `segments`.+liftSubExp :: SegLevel -> Segments -> DistInputs -> DistEnv -> SubExp -> FlattenM (Type, ResRep)+liftSubExp lvl segments inps env se = case se of+  c@(Constant prim) ->+    let t = Prim $ primValueType prim+     in ((t,) . Regular <$> letExp "lifted_const" (BasicOp $ Replicate (segmentsShape segments) c))+  Var v -> case M.lookup v $ inputReps inps env of+    Just (t, Regular v') -> do+      (t,)+        <$> case t of+          Prim {} -> pure $ Regular v'+          Array {} -> Irregular <$> mkIrregFromReg lvl segments v'+          Acc {} -> pure $ Regular v'+          Mem {} -> error "liftSubExp: Mem"+    Just (t, Irregular irreg) -> do+      irreg' <- ensureDenseIrregular lvl "lifted_irreg" irreg+      (t,)+        <$> case t of+          Prim {} -> pure $ Regular $ irregularD irreg'+          Array {} -> pure $ Irregular irreg'+          Acc {} -> error "liftSubExp: Irregular Acc"+          Mem {} -> error "liftSubExp: Mem"+    Nothing -> do+      t <- lookupType v+      v' <- letExp "free_replicated" $ BasicOp $ Replicate (segmentsShape segments) (Var v)+      (t,)+        <$> case t of+          Prim {} -> pure $ Regular v'+          Array {} -> Irregular <$> mkIrregFromReg lvl segments v'+          Acc {} -> pure $ Regular v'+          Mem {} -> error "getRepSubExp: Mem"++liftSubExpPreserveRep :: Segments -> DistInputs -> DistEnv -> SubExp -> FlattenM (Type, ResRep)+liftSubExpPreserveRep segments inps env se = case se of+  c@(Constant prim) ->+    let t = Prim $ primValueType prim+     in do+          v <- letExp "lifted_const" $ BasicOp $ Replicate (segmentsShape segments) c+          pure (t, Regular v)+  Var v -> case M.lookup v $ inputReps inps env of+    Just (t, rep) -> pure (t, rep)+    Nothing -> do+      t <- lookupType v+      v' <- letExp "free_replicated" $ BasicOp $ Replicate (segmentsShape segments) (Var v)+      pure (t, Regular v')++-- | Like 'liftSubExp' but always returns a Regular result with the+-- given expected shape. Reshapes the underlying data if necessary.+liftSubExpRegular ::+  SegLevel ->+  Segments ->+  DistInputs ->+  DistEnv ->+  Shape ->+  SubExp ->+  FlattenM VName+liftSubExpRegular lvl segments inps env expectedShape se = do+  case se of+    c@(Constant _) ->+      letExp "lifted_const" (BasicOp $ Replicate (segmentsShape segments) c)+    Var v -> liftVarRegular lvl segments inps env expectedShape v++liftVarRegular ::+  SegLevel ->+  Segments ->+  DistInputs ->+  DistEnv ->+  Shape ->+  VName ->+  FlattenM VName+liftVarRegular lvl segments inps env expectedShape x = do+  v <- case M.lookup x $ inputReps inps env of+    Just (_, Regular v') -> pure v'+    Just (_, Irregular irreg) -> do+      rep_dense <- ensureDenseIrregular lvl "lifted_irreg" irreg+      pure $ irregularD rep_dense+    Nothing ->+      letExp "free_replicated" $ BasicOp $ Replicate (segmentsShape segments) (Var x)+  v_t <- lookupType v+  if isAcc v_t || arrayShape v_t == expectedShape+    then pure v+    else+      letExp "reg_lifted" . BasicOp $+        Reshape v (reshapeAll (arrayShape v_t) expectedShape)++-- | Can this input array be lifted to a regular array? This holds unless it is+-- represented irregularly. The uniform alternatives lift their inputs regularly+-- (via 'liftSubExpRegular'), which is only valid when the inputs are actually+-- regular.+isRegularInputArr :: DistEnv -> DistInputs -> VName -> Bool+isRegularInputArr env inps arr =+  case lookup arr inps of+    Just (DistInput rt _) ->+      case resVar rt env of+        Regular {} -> True+        Irregular {} -> False+    _ -> True++liftParam :: (MonadFreshNames m) => SubExp -> FParam SOACS -> m ([FParam GPU], ResRep)+liftParam w fparam =+  case declTypeOf fparam of+    Prim pt -> do+      p <-+        newParam+          (desc <> "_lifted")+          (arrayOf (Prim pt) (Shape [w]) Nonunique)+      pure ([p], Regular $ paramName p)+    Array pt _ u -> do+      num_data <-+        newParam (desc <> "_num_data") $ Prim int64+      segments <-+        newParam (desc <> "_segments") $+          arrayOf (Prim int64) (Shape [w]) Nonunique+      flags <-+        newParam (desc <> "_F") $+          arrayOf (Prim Bool) (Shape [Var (paramName num_data)]) Nonunique+      offsets <-+        newParam (desc <> "_O") $+          arrayOf (Prim int64) (Shape [w]) Nonunique+      elems <-+        newParam (desc <> "_data") $+          arrayOf (Prim pt) (Shape [Var (paramName num_data)]) u+      pure+        ( [num_data, segments, flags, offsets, elems],+          Irregular $+            IrregularRep+              { irregularS = paramName segments,+                irregularF = paramName flags,+                irregularO = paramName offsets,+                irregularD = paramName elems,+                irregularK = Dense+              }+        )+    Acc {} ->+      error "liftParam: Acc"+    Mem {} ->+      error "liftParam: Mem"+  where+    desc = baseName (paramName fparam)++liftRegularParam :: (MonadFreshNames m) => SubExp -> FParam SOACS -> m (FParam GPU, ResRep)+liftRegularParam w fparam =+  case declTypeOf fparam of+    Prim pt -> do+      p <-+        newParam+          (desc <> "_lifted")+          (arrayOf (Prim pt) (Shape [w]) Nonunique)+      pure (p, Regular $ paramName p)+    Array pt shape u -> do+      p <-+        newParam (desc <> "_lifted") $+          arrayOf (Prim pt) (Shape [w] <> shape) u+      pure (p, Regular $ paramName p)+    Acc {} ->+      error "liftParam: Acc"+    Mem {} ->+      error "liftParam: Mem"+  where+    desc = baseName (paramName fparam)++distCerts :: DistInputs -> StmAux a -> DistEnv -> Certs+distCerts inps aux env = Certs $ map f $ unCerts $ stmAuxCerts aux+  where+    f v = case lookup v inps of+      Nothing -> v+      Just (DistInputFree vs _) -> vs+      Just (DistInput rt _) ->+        case resVar rt env of+          Regular vs -> vs+          Irregular r -> irregularD r++flattenData :: VName -> FlattenM VName+flattenData vs = do+  t <- lookupType vs+  case arrayDims t of+    [_] -> pure vs+    dims -> do+      n <- toSubExp "num_data" $ product $ map pe64 dims+      letExp (baseName vs <> "_flat") . BasicOp $+        Reshape vs $+          reshapeAll (arrayShape t) (Shape [n])++-- | Only sensible for variables of uniform type.+dataArr :: SegLevel -> Segments -> DistEnv -> DistInputs -> SubExp -> FlattenM VName+dataArr lvl _segments env inps (Var v)+  | Just v_inp <- lookup v inps =+      case v_inp of+        DistInputFree vs _ -> flattenData vs+        DistInput rt _ -> case resVar rt env of+          Irregular r -> do+            rep_dense <- ensureDenseIrregular lvl "dataArr" r+            pure $ irregularD rep_dense+          Regular vs -> flattenData vs+dataArr _ segments _ _ se = do+  -- The result is a one-dimensional array with one element per segment. With no+  -- enclosing segments there is a single implicit segment, so we replicate over+  -- a unit dimension; replicating over the empty shape would instead yield a+  -- scalar.+  let rep_shape = case segmentsShape segments of+        Shape [] -> Shape [intConst Int64 1]+        shape -> shape+  rep <- letExp "rep" $ BasicOp $ Replicate rep_shape se+  rep_t <- lookupType rep+  let dims = arrayDims rep_t+  if length dims == 1+    then pure rep+    else do+      n <- toSubExp "n" $ product $ map pe64 dims+      letExp "reshape" $ BasicOp $ Reshape rep $ reshapeAll (arrayShape rep_t) (Shape [n])++-- | Get the irregular representation of a var.+getIrregRep :: SegLevel -> Segments -> DistEnv -> DistInputs -> VName -> FlattenM IrregularRep+getIrregRep lvl segments env inps v =+  case lookup v inps of+    Just v_inp -> case v_inp of+      DistInputFree arr _ -> mkIrregFromReg lvl segments arr+      DistInput rt _ -> case resVar rt env of+        Irregular r -> pure r+        Regular arr -> mkIrregFromReg lvl segments arr+    Nothing -> do+      v' <-+        letExp (baseName v <> "_rep") . BasicOp $+          Replicate (segmentsShape segments) (Var v)+      mkIrregFromReg lvl segments v'++-- | This function walks through the *unlifted* result types+-- and uses the *lifted* results to construct the corresponding res reps.+--+-- See the 'liftResult' function for the opposite process i.e.+-- turning 'ResRep's into results.+resultToResReps :: [TypeBase s u] -> [VName] -> [ResRep]+resultToResReps types results =+  snd $+    L.mapAccumL+      ( \rs t -> case t of+          Prim {} ->+            let (v : rs') = rs+                rep = Regular v+             in (rs', rep)+          Array {} ->+            let (_ : segs : flags : offsets : elems : rs') = rs+                rep = Irregular $ IrregularRep segs flags offsets elems Dense+             in (rs', rep)+          Acc {} -> error "resultToResReps: Illegal type 'Acc'"+          Mem {} -> error "resultToResReps: Illegal type 'Mem'"+      )+      results+      types++resultToResRepsByDistResult :: [DistResult] -> [VName] -> [ResRep]+resultToResRepsByDistResult dist_res results =+  snd $+    L.mapAccumL+      ( \rs dist_res' ->+          if isRegularDistResult dist_res'+            then+              let (v : rs') = rs+               in (rs', Regular v)+            else+              let (_ : segs : flags : offsets : elems : rs') = rs+               in (rs', Irregular $ IrregularRep segs flags offsets elems Dense)+      )+      results+      dist_res++-- helper to not mess up the tags when generating new ones for the loop parameters+-- probably won't be used in future+localiseInputs :: DistEnv -> DistInputs -> FlattenM (DistInputs, DistEnv, Int)+localiseInputs env_outer inps = do+  let step (i, env_acc) (v, inp) =+        case inp of+          DistInputFree arr t ->+            pure ((i, env_acc), (v, DistInputFree arr t))+          DistInput oldrt t -> do+            let newrt = ResTag i+                rep = resVar oldrt env_outer+            env_acc' <- insertRepM newrt rep env_acc+            pure ((i + 1, env_acc'), (v, DistInput newrt t))++  ((next, env_local), inps_local) <-+    mapAccumLM step (0, mempty) inps+  pure (inps_local, env_local, next)++-- | Replicate an array to insert new inner dimensions  after the+-- existing segment dimensions.+replicateForDims :: Segments -> Shape -> VName -> FlattenM VName+replicateForDims segments dims v = do+  v_t <- lookupType v+  let seg_rank = length segments+      v_rank = arrayRank v_t+      dims_rank = shapeRank dims+      perm = [dims_rank .. dims_rank + seg_rank - 1] ++ [0 .. dims_rank - 1] ++ [seg_rank + dims_rank .. dims_rank + v_rank - 1]+  v_rep <-+    letExp (baseName v <> "_reg_rep") . BasicOp $ Replicate dims (Var v)+  letExp (baseName v <> "_reg_rep_tr") . BasicOp $ Rearrange v_rep perm++-- | Flatten a single 'DistStm', producing an updated environment.+flattenDistStm ::+  FlattenOps ->+  Segments ->+  DistEnv ->+  DistStm ->+  FlattenM DistEnv+flattenDistStm ops = flattenDistStmWith ops ops++-- | Flatten a single scalar statement, producing an updated environment.+flattenScalarStm ::+  FlattenOps ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  Stm SOACS ->+  FlattenM DistEnv+flattenScalarStm ops = flattenScalarStmAt ops (flattenSegLevel ops)++-- | 'distributeBody' with the settings in the given 'FlattenOps'.+distributeBodyWith ::+  FlattenOps ->+  Scope rep ->+  Segments ->+  DistInputs ->+  Body SOACS ->+  (DistInputs, DistStms)+distributeBodyWith ops =+  distributeBody (flattenIrregularity ops) (flattenFunHasParallelism ops)++-- | 'distributeMap' with the settings in the given 'FlattenOps'.+distributeMapWith ::+  FlattenOps ->+  Scope rep ->+  Pat Type ->+  Segments ->+  [MapArray t] ->+  Lambda SOACS ->+  (Distributed, M.Map ResTag t)+distributeMapWith ops =+  distributeMap (flattenIrregularity ops) (flattenFunHasParallelism ops)++-- | Continue flattening at the given seg level, adjusting the irregularity+-- handling mode to match. Intrablock code cannot use the machinery for+-- flattening irregular arrays, as it produces SegOps whose sizes are bound+-- inside the enclosing kernel. 'SequentialiseIrregularAll' is requested by the+-- user rather than implied by the level, so we keep it if provided.+atSegLevel :: SegLevel -> FlattenOps -> FlattenOps+atSegLevel lvl ops =+  ops {flattenSegLevel = lvl, flattenIrregularity = irreg}+  where+    irreg = case (flattenIrregularity ops, lvl) of+      (SequentialiseIrregularAll, _) -> SequentialiseIrregularAll+      (_, SegThreadInBlock {}) -> SequentialiseIrregularBasicOps+      _ -> DistributeIrregular
+ src/Futhark/Pass/Flatten/Incremental.hs view
@@ -0,0 +1,524 @@+-- | General definitions used for incremental flattening.+--+-- The idea behind incremental flattening is the observation that when+-- flattening a program+--+-- @+--   map f xs+-- @+--+-- we have two options: (i) transform @f@ to exploit any parallelism it may+-- contain, or (ii) turn @f@ into sequential code and only exploit the+-- parallelism in @map@.+--+-- In some cases we do not have a choice, e.g. if @f@ contains sufficiently+-- nonuniform operations that would result in nonuniform allocations. In other+-- cases the choice is obvious, such as when @f@ is completely scalar. However,+-- in the general case either will work, and it depends on the workload which of+-- the options is optimal: if the outer @map@ is big enough, it may be best to+-- efficiently sequentialise @f@ (which can then also permit various locality+-- optimisations, such as tiling). But if the outer @map@ does not have many+-- iterations, then we also need the parallelism in @f@ to fully saturate the+-- machine.+--+-- The idea behind incremental flattening is to generate both versions, and+-- select the appropriate one at run-time:+--+-- @+-- if predicate then sequentialise f...+--                   else parallelise f...+-- @+--+-- The predicate is based on comparing the amount of exploitable parallelism+-- with a threshold parameter. This threshold parameter is given a default value+-- based on run-time hardware characteristics, but usually has to be auto-tuned+-- in order to be optimal for a specific machine, program, and workload.+--+-- The multi-versioning approach is also used to generate more exotic versions,+-- such as one that parallelises @f@ at a deeper hardware level+-- (@Futhark.Pass.Flatten.Intrablock@).+module Futhark.Pass.Flatten.Incremental+  ( worthIntrablock,+    worthSequentialising,+    isVersionableMap,+    sufficientParallelism,+    isParallelFunInside,+    kernelAlternatives,+    intraBlockAlternative,+    mapAlternatives,+    scanRedAlternatives,+    propagateVersioningAttrs,++    -- * Transforming code+    factorScremaForParallelism,++    -- * Levels+    defaultSegLevel,+    inBlockSegLevel,+    allowVersioning,++    -- * Various queries+    bodyHasParallelism,+    lambdaHasParallelism,+    mayExploitOuter,+    onlyExploitIntra,+    mayExploitIntra,+  )+where++import Control.Monad+import Control.Monad.State+import Data.Foldable+import Data.Maybe (isJust)+import Data.Set qualified as S+import Futhark.IR.GPU+import Futhark.IR.SOACS+import Futhark.Pass.Flatten.Distribute+import Futhark.Pass.Flatten.General+import Futhark.Pass.Flatten.Intrablock qualified as Intrablock+import Futhark.Pass.Flatten.PreProcess+import Futhark.Tools+import Futhark.Transform.Rename+import Prelude hiding (div, quot, rem)++defaultSegLevel :: SegLevel+defaultSegLevel = SegThread SegVirt Nothing++inBlockSegLevel :: SegLevel+inBlockSegLevel = SegThreadInBlock SegNoVirt++allowVersioning :: SegLevel -> Bool+allowVersioning SegThreadInBlock {} = False+allowVersioning _ = True++kernelAlternatives ::+  Name ->+  [Type] ->+  Body GPU ->+  [(SubExp, Body GPU)] ->+  FlattenM [VName]+kernelAlternatives desc _ default_body [] = do+  ses <- bodyBind default_body+  forM ses $ \(SubExpRes cs se) ->+    certifying cs $+      letExp desc $+        BasicOp $+          SubExp se+kernelAlternatives desc result_ts default_body ((cond, alt) : alts) = do+  fallback_body <- do+    (fallback_vs, fallback_stms) <-+      collectStms $+        kernelAlternatives desc result_ts default_body alts+    pure $ mkBody fallback_stms $ varsRes fallback_vs++  letTupExp desc $+    Match [cond] [Case [Just $ BoolValue True] alt] fallback_body $+      MatchDec (staticShapes result_ts) MatchEquiv++cmpSizeLe ::+  Name ->+  SizeClass ->+  [SubExp] ->+  FlattenM (SubExp, Name)+cmpSizeLe desc size_class to_what = do+  x <- gets stateThresholdCounter+  modify $ \s -> s {stateThresholdCounter = x + 1}+  let size_key = desc <> "_" <> nameFromString (show x)+  to_what' <-+    letSubExp "comparatee"+      =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) to_what+  cmp_res <- letSubExp desc $ Op $ SizeOp $ CmpSizeLe size_key size_class to_what'+  pure (cmp_res, size_key)++sufficientParallelism ::+  Name ->+  [SubExp] ->+  KernelPath ->+  Maybe Int64 ->+  FlattenM (SubExp, Name)+sufficientParallelism desc ws path def =+  cmpSizeLe desc (SizeThreshold path def) ws++-- Check if the in the body there is a call to a parallel function.+-- XXX: we use this function to even reject the intra version of+-- maps that call parallel function. We should do better there.+-- One other things to note is that maybe we should create a sequential+-- version of function and replace them in these cases.+isParallelFunInside :: FunHasParallelism -> Body SOACS -> Bool+isParallelFunInside funHasParallelism = inBody+  where+    inLambda :: GLambda SOACS t -> Bool+    inLambda = any (callParallelFunction . stmExp) . bodyStms . lambdaBody+    inBody = any (callParallelFunction . stmExp) . bodyStms+    callParallelFunction (Apply fname _ _ _) = funHasParallelism fname+    callParallelFunction (BasicOp _) = False+    callParallelFunction (Match _ cases def_case _) =+      inBody def_case+        || any (inBody . caseBody) cases+    callParallelFunction (Loop _ _ body) = inBody body+    callParallelFunction (WithAcc _ lam) = inLambda lam+    callParallelFunction (Op (Stream _ _ _ lam)) = inLambda lam+    callParallelFunction (Op (Screma _ _ (ScremaForm lam _ _ _))) = inLambda lam+    callParallelFunction (Op (Hist _ _ ops lam)) =+      inLambda lam || any (inLambda . histLambda) ops+      where+        histLambda (Futhark.IR.SOACS.HistOp _ _ _ _ op) = op+    callParallelFunction (Op (FlatMap _ _ lam)) = inLambda lam+    callParallelFunction (Op JVP {}) = error "isParallelFunInside: unexpected JVP"+    callParallelFunction (Op VJP {}) = error "isParallelFunInside: unexpected VJP"+    callParallelFunction (Op WithVJP {}) = error "isParallelFunInside: unexpected WithVJP"++-- | Should we generate multiple versions for this map? This requires both that+-- we are at a level where versioning is possible ('allowVersioning') and that+-- the map itself produces only regular results (from an invariant width) and+-- does not call any parallel function (which force full flattening).+isVersionableMap :: FunHasParallelism -> SegLevel -> DistInputs -> DistEnv -> SubExp -> [DistResult] -> Lambda SOACS -> Bool+isVersionableMap funHasParallelism lvl inps _env w dist_res map_lam =+  allowVersioning lvl+    && all isRegularDistResult dist_res+    && not (isVariant inps w)+    && not (isParallelFunInside funHasParallelism (lambdaBody map_lam))++-- | Retrieve only those attributes that apply to flattening.+flatteningAttrs :: Attrs -> Attrs+flatteningAttrs = mconcat . mapAttrs p+  where+    p (AttrComp "incremental_flattening" [x]) = oneAttr x+    p (AttrComp "flattening" [x]) = oneAttr x+    p _ = mempty++onlyExploitIntra :: Attrs -> Bool+onlyExploitIntra attrs =+  "only_intra" `inAttrs` flatteningAttrs attrs++mayExploitOuter :: Attrs -> Bool+mayExploitOuter attrs =+  not $ "no_outer" `inAttrs` attrs' || "only_inner" `inAttrs` attrs'+  where+    attrs' = flatteningAttrs attrs++mayExploitIntra :: Attrs -> Bool+mayExploitIntra attrs =+  not $ "no_intra" `inAttrs` attrs' || "only_inner" `inAttrs` attrs'+  where+    attrs' = flatteningAttrs attrs++intraBlockAlternative ::+  Intrablock.IntrablockResult ->+  FlattenM (SubExp, Body GPU)+intraBlockAlternative intra = do+  addStms $ Intrablock.intraPreludeStms intra+  max_tblock_size <-+    letSubExp "max_tblock_size" $ Op $ SizeOp $ GetSizeMax SizeThreadBlock+  fits <-+    letSubExp "fits" $+      BasicOp $+        CmpOp+          (CmpSle Int64)+          (Intrablock.intraThreadBlockSize intra)+          max_tblock_size+  (intra_suff, _) <-+    sufficientParallelism+      "suff_intra_par"+      [Intrablock.intraAvailPar intra]+      mempty+      (Just Intrablock.intraMinInnerPar)+  intra_ok <-+    letSubExp "intra_suff_and_fits" $+      BasicOp $+        BinOp LogAnd fits intra_suff+  intra_body <-+    renameBody $+      mkBody+        (Intrablock.intraKernelStms intra)+        (varsRes $ Intrablock.intraResultNames intra)+  pure (intra_ok, intra_body)++-- | Construct the multi-versioned alternatives for a map, given the+-- fully-flattened body, the outer-parallel-only body, and an optional+-- intrablock result. This is the shared versioning policy used both for+-- top-level maps and for maps nested inside a map-nest; the only differences+-- between the two are which bodies are supplied and how their results are+-- consumed, both of which are handled by the caller. The @ws@ are the widths+-- whose product bounds the outer parallelism (used for the threshold+-- comparison). Returns the names bound to the final results.+mapAlternatives ::+  -- | Description for the result bindings.+  Name ->+  [Type] ->+  Attrs ->+  -- | Does the map body call a parallel function? If so we must fully flatten.+  Bool ->+  -- | Is the body worth sequentialising (offering an outer-only version)?+  Bool ->+  [SubExp] ->+  Body GPU ->+  Body GPU ->+  Maybe Intrablock.IntrablockResult ->+  FlattenM [VName]+mapAlternatives desc result_ts attrs parallel_fun_inside worth_seq ws full_body outer_body intra' =+  case intra' of+    _+      | parallel_fun_inside ->+          kernelAlternatives desc result_ts full_body []+      | "sequential_inner" `inAttrs` attrs ->+          kernelAlternatives desc result_ts outer_body []+    Nothing+      | not only_intra,+        worth_seq,+        mayExploitOuter attrs -> do+          (outer_suff, _) <- outerSuff+          kernelAlternatives desc result_ts full_body [(outer_suff, outer_body)]+      | otherwise ->+          kernelAlternatives desc result_ts full_body []+    Just intra_res+      | only_intra -> do+          (_, intra_body) <- intraBlockAlternative intra_res+          kernelAlternatives desc result_ts intra_body []+      | worth_seq,+        mayExploitOuter attrs -> do+          (outer_suff, _) <- outerSuff+          intra_alt <- intraBlockAlternative intra_res+          kernelAlternatives desc result_ts full_body [(outer_suff, outer_body), intra_alt]+      | otherwise -> do+          intra_alt <- intraBlockAlternative intra_res+          kernelAlternatives desc result_ts full_body [intra_alt]+  where+    only_intra = onlyExploitIntra attrs++    outerSuff = sufficientParallelism suffOuterPar ws mempty Nothing++-- | Construct the multi-versioned alternatives for a scan or reduce, given the+-- fully-flattened body and the outer-parallel-only body. Unlike+-- 'mapAlternatives' there is no intrablock version, and the outer-only version+-- is always offered (subject to attributes). Shared between top-level and+-- nested uniform scans/reduces.+scanRedAlternatives ::+  Name ->+  [Type] ->+  Attrs ->+  -- | Does the operator body call a parallel function? If so we must fully flatten.+  Bool ->+  -- | Does the seg level permit versioning at all (false in-block)?+  Bool ->+  [SubExp] ->+  Body GPU ->+  Body GPU ->+  FlattenM [VName]+scanRedAlternatives desc result_ts attrs parallel_fun_inside allow_versioning ws full_body outer_body+  | parallel_fun_inside =+      fullAlternative+  | "sequential_inner" `inAttrs` attrs =+      outerAlternative+  | mayExploitOuter attrs && allow_versioning =+      fullWithOuterAlternative+  | otherwise =+      fullAlternative+  where+    fullAlternative = kernelAlternatives desc result_ts full_body []++    outerAlternative = kernelAlternatives desc result_ts outer_body []++    fullWithOuterAlternative = do+      (outer_suff, _) <- sufficientParallelism suffOuterPar ws mempty Nothing+      kernelAlternatives desc result_ts full_body [(outer_suff, outer_body)]++-- | The name of the threshold parameter that is used to select outer-only+-- parallelism.+suffOuterPar :: Name+suffOuterPar = "suff_outer_par"++-- | Intra-group parallelism is worthwhile if the lambda contains more+-- than one instance of non-map nested parallelism, or any nested+-- parallelism inside a loop.+worthIntrablock :: Lambda SOACS -> Bool+worthIntrablock lam =+  bodyInterest (lambdaBody lam) > 1+  where+    bodyInterest body =+      sum $ interest <$> bodyStms body+    interest stm+      | "sequential" `inAttrs` attrs =+          0 :: Int+      | Op (Screma w _ form) <- stmExp stm,+        Just lam' <- isMapSOAC form =+          mapLike w lam'+      | Loop _ _ body <- stmExp stm =+          bodyInterest body * 10+      | Match _ cases defbody _ <- stmExp stm =+          foldl+            max+            (bodyInterest defbody)+            (map (bodyInterest . caseBody) cases)+      | Op (Screma w _ (ScremaForm lam' _ _ _)) <- stmExp stm =+          zeroIfTooSmall w + bodyInterest (lambdaBody lam')+      | Op (Stream _ _ _ lam') <- stmExp stm =+          bodyInterest $ lambdaBody lam'+      | WithAcc _ lam' <- stmExp stm =+          bodyInterest $ lambdaBody lam'+      | otherwise =+          0+      where+        attrs = stmAuxAttrs $ stmAux stm+        sequential_inner = "sequential_inner" `inAttrs` attrs++        zeroIfTooSmall (Constant (IntValue x))+          | intToInt64 x < 32 = 0+        zeroIfTooSmall _ = 1++        mapLike w lam' =+          if sequential_inner+            then 0+            else max (zeroIfTooSmall w) (bodyInterest (lambdaBody lam'))++-- | A lambda is worth sequentialising if it contains enough nested parallelism+-- of an interesting kind, or if distributing it would fragment sequential+-- control flow - that is, if it contains meaningful parallelism nested inside a+-- sequential loop or branch. Distribution must then split the loop or branch+-- into separate kernel launches (and possibly host-evaluated control flow) per+-- sequential step, so a version that instead sequentialises the nested+-- parallelism is always worth offering.+worthSequentialising :: Lambda SOACS -> Bool+worthSequentialising lam =+  bodyInterest (0 :: Int) (lambdaBody lam) > 1+  where+    bodyInterest depth body =+      sum $ interest depth <$> bodyStms body+    interest depth stm+      | "sequential" `inAttrs` attrs =+          0 :: Int+      | Op (Screma _ _ form@(ScremaForm lam' _ _ _)) <- stmExp stm,+        isJust $ isMapSOAC form =+          if sequential_inner+            then 0+            else bodyInterest (depth + 1) (lambdaBody lam')+      | Loop _ _ body <- stmExp stm =+          bodyInterest (depth + 1) body * 10+      | Match _ cases defbody _ <- stmExp stm =+          (2 *) $+            maximum $+              map (bodyInterest (depth + 1)) $+                defbody : map caseBody cases+      | WithAcc _ withacc_lam <- stmExp stm =+          bodyInterest (depth + 1) (lambdaBody withacc_lam)+      | Op (Screma _ _ form@(ScremaForm lam' _ _ _)) <- stmExp stm =+          1+            + bodyInterest (depth + 1) (lambdaBody lam')+            ++            -- Give this a bigger score if it's a redomap just inside+            -- the the outer lambda, as these are often tileable and+            -- thus benefit more from sequentialisation.+            case (isRedomapSOAC form, depth) of+              (Just _, 0) -> 1+              _ -> 0+      | Op (Stream _ _ _ lam') <- stmExp stm =+          bodyInterest (depth + 1) (lambdaBody lam')+      | otherwise =+          0+      where+        attrs = stmAuxAttrs $ stmAux stm+        sequential_inner = "sequential_inner" `inAttrs` attrs++bodyHasParallelism :: FunHasParallelism -> Body SOACS -> Bool+bodyHasParallelism funHasParallelism =+  any (isParallelStm funHasParallelism) . bodyStms++lambdaHasParallelism :: FunHasParallelism -> Lambda SOACS -> Bool+lambdaHasParallelism funHasParallelism =+  bodyHasParallelism funHasParallelism . lambdaBody++-- | Like 'lambdaHasParallelism', but only counts meaningful+-- parallelism: a SOAC, a call to a parallel function, or a statement+-- with an irregular result, which requires flattening to exploit.+-- Basic operations such as 'Replicate' of invariant size do not+-- provide enough parallelism on their own to make multi-versioning+-- worthwhile.  See Note [Meaningful Parallelism] in+-- Futhark.Pass.Flatten.Distribute.+lambdaHasMeaningfulParallelism :: FunHasParallelism -> Lambda SOACS -> Bool+lambdaHasMeaningfulParallelism funHasParallelism lam =+  any interesting $ bodyStms $ lambdaBody lam+  where+    free_in_lam = freeIn lam+    invariantDim (Var v) = v `nameIn` free_in_lam+    invariantDim Constant {} = True+    irregularResult =+      not . all (all invariantDim . arrayDims) . patTypes . stmPat+    interesting stm =+      stmHasMeaningfulParallelism funHasParallelism stm || irregularResult stm++-- | Produce a body suitable for full flattening from a Screma, or+-- 'Nothing' if none of its lambdas contain meaningful parallelism, in+-- which case multi-versioning is not worthwhile.  See Note+-- [Meaningful Parallelism] in Futhark.Pass.Flatten.Distribute.+factorScremaForParallelism ::+  (MonadBuilder m) =>+  FunHasParallelism ->+  Scope SOACS ->+  Certs ->+  Pat Type ->+  SubExp ->+  [VName] ->+  ScremaForm SOACS ->+  m (Maybe (Body SOACS))+factorScremaForParallelism funHasParallelism scope certs pat w arrs form+  | Just (reds, map_lam) <- isRedomapSOAC form,+    lambdaHasMeaningfulParallelism funHasParallelism map_lam = do+      map_lam' <- preprocessLambda scope map_lam+      (map_stm, red_stm) <-+        redomapToMapAndReduce+          pat+          (w, reds, map_lam', arrs)+      Just <$> mkFactoredBody (stmsFromList [map_stm, red_stm])+  | Just (post_lam, scans, map_lam) <- isMaposcanomapSOAC form,+    lambdaHasMeaningfulParallelism funHasParallelism map_lam,+    lambdaHasMeaningfulParallelism funHasParallelism post_lam = do+      map_lam' <- preprocessLambda scope map_lam+      post_lam' <- preprocessLambda scope post_lam+      (map_stm, scan_stm, post_stm) <-+        maposcanomapToMapScanAndMap+          pat+          (w, post_lam', scans, map_lam', arrs)+      Just <$> mkFactoredBody (stmsFromList [map_stm, scan_stm, post_stm])+  | Just (post_lam, scans, map_lam) <- isMaposcanomapSOAC form,+    lambdaHasMeaningfulParallelism funHasParallelism map_lam = do+      map_lam' <- preprocessLambda scope map_lam+      post_lam' <- preprocessLambda scope post_lam+      (map_stm, scanomap_stm) <-+        maposcanomapToMaposcanAndMap+          pat+          (w, post_lam', scans, map_lam', arrs)+      Just <$> mkFactoredBody (stmsFromList [map_stm, scanomap_stm])+  | Just (post_lam, scans, map_lam) <- isMaposcanomapSOAC form,+    lambdaHasMeaningfulParallelism funHasParallelism post_lam = do+      map_lam' <- preprocessLambda scope map_lam+      post_lam' <- preprocessLambda scope post_lam+      (map_stm, scan_stm, post_stm) <-+        maposcanomapToMapScanAndMap+          pat+          (w, post_lam', scans, map_lam', arrs)+      Just <$> mkFactoredBody (stmsFromList [map_stm, scan_stm, post_stm])+  | otherwise =+      pure Nothing+  where+    mkFactoredBody stms = do+      stms' <- fmap (certify certs) <$> preprocessStms scope stms+      pure $ mkBody stms' $ varsRes $ patNames pat++-- | Propagate incremental flattening attributes to the statements of+-- a map lambda body. Statements that carry their own incremental+-- flattening attributes are left alone.+propagateVersioningAttrs :: Attrs -> Lambda SOACS -> Lambda SOACS+propagateVersioningAttrs attrs lam+  | attrs' == mempty = lam+  | otherwise =+      lam {lambdaBody = (lambdaBody lam) {bodyStms = fmap onStm (bodyStms (lambdaBody lam))}}+  where+    attrs' = versioningAttrs attrs+    onStm stm+      | versioningAttrs (stmAuxAttrs (stmAux stm)) == mempty =+          stm {stmAux = (stmAux stm) {stmAuxAttrs = attrs' <> stmAuxAttrs (stmAux stm)}}+      | otherwise = stm+    versioningAttrs (Attrs s) = Attrs $ S.filter isVersioningAttr s+    isVersioningAttr (AttrComp "incremental_flattening" _) = True+    isVersioningAttr (AttrComp "flattening" _) = True+    isVersioningAttr _ = False
+ src/Futhark/Pass/Flatten/Intrablock.hs view
@@ -0,0 +1,353 @@+{-# LANGUAGE TypeFamilies #-}++module Futhark.Pass.Flatten.Intrablock+  ( IntrablockResult (..),+    intrablockParallelise,+    intraMinInnerPar,+    InBlockMapTransformer,+  )+where++import Control.Monad+import Control.Monad.RWS+import Control.Monad.Writer+import Data.Map qualified as M+import Data.Set qualified as S+import Futhark.Analysis.PrimExp.Convert+import Futhark.IR.GPU hiding (HistOp)+import Futhark.IR.SOACS+import Futhark.MonadFreshNames+import Futhark.Pass.Flatten.Distribute+import Futhark.Pass.Flatten.General+import Futhark.Pass.Flatten.PreProcess (preprocessLambda)+import Futhark.Tools+import Futhark.Transform.FirstOrderTransform qualified as FOT+import Futhark.Transform.Rename+import Futhark.Transform.ToGPU+import Prelude hiding (log)++-- | The minimum amount of inner parallelism we require (by default)+-- in intra-group versions.+intraMinInnerPar :: Int64+intraMinInnerPar = 32++data IntrablockResult = IntrablockResult+  { intraMinPar :: SubExp,+    intraAvailPar :: SubExp,+    intraThreadBlockSize :: SubExp,+    intraPreludeStms :: Stms GPU,+    intraKernelStms :: Stms GPU,+    intraResultNames :: [VName]+  }++type InBlockMapTransformer =+  Pat Type ->+  SubExp ->+  [VName] ->+  Lambda SOACS ->+  FlattenM ()++foldBinOp' :: (MonadBuilder m) => BinOp -> [SubExp] -> m (Exp (Rep m))+foldBinOp' _ [] = eSubExp $ intConst Int64 1+foldBinOp' bop (x : xs) = foldBinOp bop x xs++-- | Extract the intra-block parallelism used by the body.+findParallelism :: KernelBody GPU -> [[SubExp]]+findParallelism = S.toList . execWriter . onKernelBody+  where+    onKernelBody = mapM_ onStm . bodyStms+    onBody = mapM_ onStm . bodyStms+    onStm stm = walkExpM walker $ stmExp stm+    walker =+      (identityWalker @GPU)+        { walkOnBody = const onBody,+          walkOnOp = onOp+        }+    onOp (SegOp op) = do+      tell $ S.singleton $ segSpaceDims $ segSpace op+      onKernelBody $ segBody op+    onOp _ = pure ()++computeThreadBlockSize :: [[SubExp]] -> [[SubExp]] -> FlattenM (SubExp, SubExp)+computeThreadBlockSize wss_min wss_avail = do+  ws_min <-+    mapM (letSubExp "one_intra_par_min" <=< foldBinOp' (Mul Int64 OverflowUndef)) $+      filter (not . null) wss_min+  ws_avail <-+    mapM (letSubExp "one_intra_par_avail" <=< foldBinOp' (Mul Int64 OverflowUndef)) $+      filter (not . null) wss_avail++  -- The amount of parallelism available *in the worst case* is+  -- equal to the smallest parallel loop, or *at least* 1.+  intra_avail_par <-+    letSubExp "intra_avail_par" =<< foldBinOp' (SMin Int64) ws_avail++  tblock_size <- newVName "computed_tblock_size"+  -- The group size is either the maximum of the minimum parallelism+  -- exploited, or the desired parallelism (bounded by the max group+  -- size) in case there is no minimum.+  letBindNames [tblock_size]+    =<< if null ws_min+      then+        eBinOp+          (SMin Int64)+          (eSubExp =<< letSubExp "max_tblock_size" (Op $ SizeOp $ GetSizeMax SizeThreadBlock))+          (eSubExp intra_avail_par)+      else foldBinOp' (SMax Int64) ws_min+  pure (intra_avail_par, Var tblock_size)++-- | Check whether this result is actually something that is acceptable to use:+-- the parallel dimensions must all be bound outside the kernel, as nonuniform+-- parallelism cannot be exploited inside a thread block.+noNonuniformPar :: Names -> FlattenM IntrablockResult -> FlattenM (Maybe IntrablockResult)+noNonuniformPar pars m = do+  outside_scope <- askScope+  if allNames (`M.member` outside_scope) pars+    then Just <$> m+    else pure Nothing++intrablockParallelise ::+  InBlockMapTransformer ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  Pat Type ->+  StmAux () ->+  SubExp ->+  [VName] ->+  Lambda SOACS ->+  FlattenM (Maybe IntrablockResult)+intrablockParallelise map_in_block segments env inps dist_res _pat aux w arrs lam0 = do+  gpu_scope <- askScope+  let pp_scope = castScope $ scopeOfDistInputs inps <> gpu_scope+  lam <- renameLambda =<< preprocessLambda pp_scope lam0++  let result_ts =+        [ t `arrayOfShape` segmentsShape segments+        | DistResult _ (DistType _ _ t) _ <- dist_res+        ]++  -- Reconstruct the per-enclosing-segment inputs, so they can be read at the+  -- appropriate indices inside the kernel.+  (param_inputs, input_prelude_stms) <-+    collectStms $+      zipWithM (prepareRegularMapInput segments env inps) (lambdaParams lam) arrs++  -- A free variable of the map body is a per-enclosing-segment value: it is+  -- brought into the body's scope as a distributed input and reconstructed at+  -- the outer indices inside the kernel (via 'readInBlockInputs' below). With no+  -- enclosing segments, however, it is an ordinary top-level value already in+  -- the kernel's outer scope, so there is nothing to reconstruct - hence the+  -- empty 'free_inputs', which makes both the scope extension and the read below+  -- vanish.+  free_inputs <- if null segments then pure [] else freeInputsFor inps lam+  kbody <-+    localScope (scopeOfDistInputs free_inputs <> scopeOfLParams (lambdaParams lam)) $+      intrablockParalleliseBody map_in_block (lambdaBody lam)++  nested_pat <-+    fmap Pat $+      zipWith PatElem+        <$> mapM (newName . distResName) dist_res+        <*> pure result_ts++  let new_segments = segments <> pure w+      wss = findParallelism kbody++  noNonuniformPar (freeIn wss) $ do+    ((intra_avail_par, tblock_size, kspace, num_tblocks), prelude_stms) <-+      collectStms $ do+        num_tblocks <-+          letSubExp "intra_num_tblocks"+            =<< foldBinOp' (Mul Int64 OverflowUndef) new_segments+        (intra_avail_par, tblock_size) <- computeThreadBlockSize wss wss+        gtids <- mapM (const $ newVName "gtid") new_segments+        kspace <- mkSegSpace $ zip gtids new_segments+        pure (intra_avail_par, tblock_size, kspace, num_tblocks)++    read_input_stms <-+      collectStms_ . localScope (scopeOfSegSpace kspace <> scopeOf input_prelude_stms <> scopeOf prelude_stms) $ do+        let SegSpace _ gtids_and_dims = kspace+            full_is = map (Var . fst) gtids_and_dims+            outer_is = take (segmentsRank segments) full_is+        readInBlockInputs segments env outer_is free_inputs+        readInBlockInputs new_segments mempty full_is param_inputs++    let kbody' = kbody {bodyStms = read_input_stms <> bodyStms kbody}+        rts = map (length new_segments `stripArray`) result_ts+        grid = KernelGrid (Count num_tblocks) (Count tblock_size)+        lvl = SegBlock SegNoVirt (Just grid)+        kstm = Let nested_pat aux $ Op $ SegOp $ SegMap lvl kspace rts kbody'++    pure $+      IntrablockResult+        { intraMinPar = intra_avail_par,+          intraAvailPar = intra_avail_par,+          intraThreadBlockSize = tblock_size,+          intraPreludeStms = input_prelude_stms <> prelude_stms,+          intraKernelStms = oneStm kstm,+          intraResultNames = patNames nested_pat+        }++readInBlockInputs :: Segments -> DistEnv -> [SubExp] -> DistInputs -> FlattenM ()+readInBlockInputs segments env is inputs =+  mapM_ onInput inputs+  where+    onInput (v, inp) = do+      v' <- readInputVar segments env is inputs v+      let t = distInputType inp+      if isAcc t+        then+          letBindNames [v] $ BasicOp $ SubExp $ Var v'+        else+          if arrayRank t > 0+            then+              letBindNames [v] $ BasicOp $ Replicate mempty $ Var v'+            else+              letBindNames [v] $ BasicOp $ SubExp $ Var v'++prepareRegularMapInput ::+  Segments ->+  DistEnv ->+  DistInputs ->+  Param Type ->+  VName ->+  FlattenM (VName, DistInput)+prepareRegularMapInput segments env inps p arr = do+  t <- lookupInputType inps arr+  let expectedShape = segmentsShape segments <> arrayShape t+      lvl = SegThread SegVirt Nothing+  arr_rep <- liftVarRegular lvl segments inps env expectedShape arr+  pure (paramName p, DistInputFree arr_rep (paramType p))++freeInputsFor :: DistInputs -> Lambda SOACS -> FlattenM DistInputs+freeInputsFor inps lam =+  do+    let free = freeIn lam+    free_sizes <-+      foldMap freeIn <$> mapM (lookupInputType inps) (namesToList free)+    pure+      [ (v, inp)+      | v <- namesToList $ free <> free_sizes,+        Just inp <- [lookup v inps]+      ]++intrablockBody :: InBlockMapTransformer -> Body SOACS -> FlattenM (Body GPU)+intrablockBody map_in_block body = do+  stms <- collectStms_ $ intrablockStms map_in_block $ bodyStms body+  pure $ mkBody stms $ bodyResult body++intrablockLambda :: InBlockMapTransformer -> Lambda SOACS -> FlattenM (Lambda GPU)+intrablockLambda map_in_block lam =+  mkLambda (lambdaParams lam) $+    bodyBind =<< intrablockBody map_in_block (lambdaBody lam)++intrablockWithAccInput :: InBlockMapTransformer -> WithAccInput SOACS -> FlattenM (WithAccInput GPU)+intrablockWithAccInput _ (shape, arrs, Nothing) =+  pure (shape, arrs, Nothing)+intrablockWithAccInput map_in_block (shape, arrs, Just (lam, nes)) = do+  lam' <- intrablockLambda map_in_block lam+  pure (shape, arrs, Just (lam', nes))++intrablockStm :: InBlockMapTransformer -> Stm SOACS -> FlattenM ()+intrablockStm map_in_block stm@(Let pat aux e) = do+  scope <- askScope+  let lvl = SegThreadInBlock SegNoVirt++  case e of+    Loop merge form loopbody ->+      localScope (scopeOfLoopForm form <> scopeOfFParams (map fst merge)) $ do+        loopbody' <- intrablockBody map_in_block loopbody+        certifying (stmAuxCerts aux) . letBind pat $+          Loop merge form loopbody'+    Match cond cases defbody ifdec -> do+      cases' <- mapM (traverse $ intrablockBody map_in_block) cases+      defbody' <- intrablockBody map_in_block defbody+      certifying (stmAuxCerts aux) . letBind pat $+        Match cond cases' defbody' ifdec+    WithAcc inputs lam -> do+      inputs' <- mapM (intrablockWithAccInput map_in_block) inputs+      lam' <- intrablockLambda map_in_block lam+      certifying (stmAuxCerts aux) . letBind pat $ WithAcc inputs' lam'+    Op soac+      | "sequential_outer" `inAttrs` stmAuxAttrs aux ->+          intrablockStms map_in_block . fmap (certify (stmAuxCerts aux))+            =<< runBuilder_ (FOT.transformSOAC pat soac)+    Op (Screma w arrs form)+      | Just lam <- isMapSOAC form ->+          map_in_block pat w arrs lam+    Op (Screma w arrs form)+      | Just (post_lam, scans, mapfun) <- isMaposcanomapSOAC form,+        -- FIXME: Futhark.CodeGen.ImpGen.GPU.Block.compileGroupOp+        -- cannot handle multiple scan operators yet.+        Scan scanfun nes <- singleScan scans -> do+          let scanfun' = soacsLambdaToGPU scanfun+              mapfun' = soacsLambdaToGPU mapfun+              post_op = soacsLambdaToGPU post_lam+          scan_res <- certifying (stmAuxCerts aux) $ genUniformSegScanomapWithPost lvl (pure w) "intra_maposcanomap" scanfun' mempty nes post_op mapfun' arrs (const $ pure ())+          zipWithM_+            ( \pe v ->+                letBindNames [patElemName pe] $ BasicOp $ SubExp $ Var v+            )+            (patElems pat)+            scan_res+    Op (Screma w arrs form)+      | Just (reds, map_lam) <- isRedomapSOAC form -> do+          let onRed red =+                let red_lam = redLambda red+                    comm+                      | commutativeLambda red_lam = Commutative+                      | otherwise = redComm red+                 in Reduce comm (soacsLambdaToGPU red_lam) (redNeutral red)+              reds_gpu = map onRed reds+              map_lam' = soacsLambdaToGPU map_lam+          (red_res, stms) <- runBuilder (genUniformSegRed lvl "intra_redomap" (pure w) reds_gpu mempty map_lam' arrs (const $ pure ()))+          certifying (stmAuxCerts aux) $ do+            addStms stms+            zipWithM_+              ( \pe v ->+                  letBindNames [patElemName pe] $ BasicOp $ SubExp $ Var v+              )+              (patElems pat)+              red_res+    Op (Screma w arrs form) ->+      -- This screma is too complicated for us to immediately do+      -- anything, so split it up and try again.+      mapM_ (intrablockStm map_in_block) . fmap (certify (stmAuxCerts aux)) . snd+        =<< runBuilderT (dissectScrema pat w form arrs) (scopeForSOACs scope)+    Op (Hist w arrs ops bucket_fun) -> do+      let bucket_fun' = soacsLambdaToGPU bucket_fun++      (hist_res, stms) <- runBuilder (genUniformSegHist lvl "Uniform_segHist" (pure w) ops bucket_fun' arrs (const $ pure ()))+      certifying (stmAuxCerts aux) $ do+        addStms stms+        zipWithM_+          ( \pe v ->+              letBindNames [patElemName pe] $+                BasicOp $+                  SubExp $+                    Var v+          )+          (patElems pat)+          hist_res+    Op (Stream w arrs accs lam) -> do+      types <- asksScope castScope+      ((), stream_stms) <-+        runBuilderT (sequentialStreamWholeArray pat w accs lam arrs) types+      intrablockStms map_in_block stream_stms+    _ ->+      addStm $ soacsStmToGPU stm++intrablockStms :: InBlockMapTransformer -> Stms SOACS -> FlattenM ()+intrablockStms map_in_block = mapM_ $ intrablockStm map_in_block++intrablockParalleliseBody ::+  InBlockMapTransformer ->+  Body SOACS ->+  FlattenM (KernelBody GPU)+intrablockParalleliseBody map_in_block body = do+  kstms <- collectStms_ $ intrablockStms map_in_block $ bodyStms body+  pure $ Body () kstms $ map ret $ bodyResult body+  where+    ret (SubExpRes cs se) = Returns ResultMaySimplify cs se
+ src/Futhark/Pass/Flatten/Loop.hs view
@@ -0,0 +1,635 @@+module Futhark.Pass.Flatten.Loop+  ( flattenLoop,+  )+where++import Control.Monad+import Control.Monad.Reader (runReaderT)+import Data.Containers.ListUtils (nubOrd)+import Data.Foldable+import Data.List qualified as L+import Data.Map qualified as M+import Data.Set qualified as S+import Data.Tuple.Solo+import Futhark.IR.GPU+import Futhark.IR.SOACS+import Futhark.IR.SOACS.Simplify (simplifyStms)+import Futhark.MonadFreshNames+import Futhark.Pass.Flatten.Distribute+import Futhark.Pass.Flatten.General+import Futhark.Tools+import Prelude hiding (div, quot, rem)++-- | Is this dimension variant to the loop or the outer map context -+-- either because it is itself a loop parameter, or because it is+-- variant in the outer map-nest?+variantDim :: DistInputs -> S.Set VName -> SubExp -> Bool+variantDim _ _ Constant {} = False+variantDim inps loopParamNames (Var v) =+  v `S.member` loopParamNames || isVariant inps (Var v)++-- Check whether a loop parameter array needs irregular representation.+-- we need the irregular representation when any of its dimensions are either:+-- a loop parameter name or variant in the outer map context+needsIrregular :: DistInputs -> S.Set VName -> DeclType -> Bool+needsIrregular inps loopParamNames t =+  case t of+    Array {} -> any (variantDim inps loopParamNames) (arrayDims t)+    _ -> False++-- Lift a loop parameter and its initial value together.+-- If the parameter is an array whose dimensions are all invariant,+-- we lift it to a regular array. Otherwise we fall back to irregular.+liftLoopParam ::+  SegLevel ->+  Segments ->+  SubExp ->+  DistInputs ->+  DistEnv ->+  S.Set VName ->+  (FParam SOACS, SubExp) ->+  FlattenM ([FParam GPU], ResRep, [SubExp])+liftLoopParam lvl segments num_segments inps env loopParamNames (fparam, initSE) = do+  let t = declTypeOf fparam+  case t of+    Prim pt -> do+      param <-+        newParam+          (baseName (paramName fparam) <> "_lifted")+          (arrayOf (Prim pt) (segmentsShape segments) Nonunique)+      initV <- liftSubExpRegular lvl segments inps env (segmentsShape segments) initSE+      pure ([param], Regular $ paramName param, [Var initV])+    Array pt _ u+      | needsIrregular inps loopParamNames t -> do+          (params, rep) <- liftParam num_segments fparam+          (_, initRep) <- liftSubExp lvl segments inps env initSE+          irreg <- case initRep of+            -- This will not happen.+            Regular v -> mkIrregFromReg lvl segments v+            Irregular irreg -> pure irreg+          initVals <- irregularRepToFlatArrs num_segments irreg+          pure (params, rep, map Var initVals)+      | otherwise -> do+          -- Regular case: all dims are invariant, just add w as outermost dim+          let pShape = segmentsShape segments <> arrayShape t+          p <-+            newParam+              (baseName (paramName fparam) <> "_lifted")+              (arrayOf (Prim pt) pShape u)+          initV <- liftSubExpRegular lvl segments inps env pShape initSE+          -- If the parameter is consumed, we must not consume the+          -- representation array (it may be used by other versions in+          -- multi-versioned code), so insert a copy. The simplifier hopefully+          -- removes it again when consuming the representation directly is+          -- safe.+          initV' <-+            if u == Unique+              then letExp (baseName (paramName fparam) <> "_inter_copy") =<< eCopy (eVar initV)+              else pure initV+          pure ([p], Regular $ paramName p, [Var initV'])+    Acc {} -> do+      initV <- liftSubExpRegular lvl segments inps env mempty initSE+      let Param attrs v acc_t = fparam+      param <- Param attrs <$> newName v <*> pure acc_t+      pure ([param], Regular $ paramName param, [Var initV])+    Mem {} ->+      error "liftLoopParam: Mem"++-- | Construct the body of an interchanged uniform loop: a single+-- Screma mapping the original loop body over the lifted loop+-- parameters (and any other inputs used by the body), transformed as+-- if it were a top-level statement - in particular, it is subject to+-- multi-versioning. The lambda parameters reuse the original names,+-- so the body can be used unchanged. Only usable when all involved+-- values are regular.+interchangedLoopBody ::+  FlattenOps ->+  SubExp ->+  Segments ->+  DistEnv ->+  [(FParam SOACS, FParam GPU)] ->+  DistInputs ->+  StmAux () ->+  Body SOACS ->+  FlattenM (Body GPU)+interchangedLoopBody ops num_segments segments env params free_inps aux body = buildBody_ $ do+  let flatInput name arr t = do+        arr_t <- lookupType arr+        letExp (baseName name <> "_flat") . BasicOp . Reshape arr $+          reshapeAll (arrayShape arr_t) (Shape [num_segments] <> arrayShape t)+      inputArr (DistInputFree arr _) = arr+      inputArr (DistInput rt _) = case resVar rt env of+        Regular arr -> arr+        Irregular {} -> error "interchangedLoopBody: irregular input"+  param_arrs <- forM params $ \(p, lifted_p) ->+    flatInput (paramName p) (paramName lifted_p) (fromDecl (declTypeOf p))+  free_arrs <- forM free_inps $ \(v, inp) ->+    flatInput v (inputArr inp) (distInputType inp)++  let lam_params =+        [Param mempty (paramName p) (fromDecl (declTypeOf p)) | (p, _) <- params]+          ++ [Param mempty v (distInputType inp) | (v, inp) <- free_inps]+      row_ts = [fromDecl (declTypeOf p) | (p, _) <- params]+      lam = Lambda lam_params row_ts body+  pes <- forM (zip params row_ts) $ \((p, _), t) ->+    PatElem+      <$> newName (paramName p)+      <*> pure (t `arrayOfRow` num_segments)+  form <- mapSOAC lam+  let map_stm :: Stm SOACS+      map_stm =+        Let (Pat pes) (aux {stmAuxCerts = mempty}) $+          Op $+            Screma num_segments (param_arrs ++ free_arrs) form++  -- Simplify before transforming. Apart from generally producing+  -- better code, this hoists statements that are invariant to the+  -- mapped values out of the Screma, and in particular any sizes+  -- they compute must be in scope when the Screma is versioned+  -- (e.g. for deciding intrablock feasibility).+  scope <- castScope <$> askScope+  map_stms <- runReaderT (simplifyStms (oneStm map_stm)) (scope :: Scope SOACS)+  mapM_ (flattenTopLevelStm ops) map_stms+  fmap (map (SubExpRes mempty . Var)) . forM (zip pes params) $ \(pe, (p, _)) -> do+    pe_t <- lookupType (patElemName pe)+    let seg_shape = segmentsShape segments <> arrayShape (fromDecl (declTypeOf p))+    letExp (baseName (paramName p) <> "_unflat") . BasicOp . Reshape (patElemName pe) $+      reshapeAll (arrayShape pe_t) seg_shape++liftLoopResult :: SegLevel -> Segments -> SubExp -> DistInputs -> DistEnv -> DistResult -> SubExpRes -> FlattenM Result+liftLoopResult lvl segments num_segments inps env dist_res res =+  if isRegularDistResult dist_res+    then do+      let (DistType _ _ t) = distResType dist_res+      let expectedShape = segmentsShape segments <> arrayShape t+      v <- liftSubExpRegular lvl segments inps env expectedShape (resSubExp res)+      pure [SubExpRes mempty (Var v)]+    else case resSubExp res of+      Var v -> do+        irreg <- getIrregRep lvl segments env inps v+        varsRes <$> irregularRepToFlatArrs num_segments irreg+      _ -> error "liftLoopResult: irregular result is not a variable"++-- | Distribute the loop body statement by statement and lift the+-- distributed statements, producing the statements and result of the+-- body of the lifted loop. The provided scope is that of the lifted+-- loop parameters (and any loop index); it is brought into scope only+-- after distribution, as the original body cannot reference it.+distributedLoopBody ::+  FlattenOps ->+  Segments ->+  SubExp ->+  Scope GPU ->+  DistInputs ->+  DistEnv ->+  [DistResult] ->+  Body SOACS ->+  FlattenM (Body GPU)+distributedLoopBody ops segments num_segments loop_scope inputs env res body = do+  scope <- askScope+  let lvl = flattenSegLevel ops+      (inputs', dstms) =+        distributeBodyWith ops scope segments inputs body+  buildBody_ $ localScope loop_scope $ do+    env' <- foldM (flattenDistStm ops segments) env dstms+    concat <$> zipWithM (liftLoopResult lvl segments num_segments inputs' env') res (bodyResult body)++-- | Make the original loop parameters available as distribution+-- inputs for the loop body, mapped to their lifted representations.+loopBodyInputs :: DistEnv -> DistInputs -> [FParam SOACS] -> [ResRep] -> FlattenM (DistInputs, DistEnv)+loopBodyInputs env inps old_loop_params lifted_loop_reps = do+  (inps_local, env_local, next) <- localiseInputs env inps+  let loop_param_inputs =+        zipWith+          (\p j -> (paramName p, DistInput (ResTag j) (paramType p)))+          old_loop_params+          [next ..]+      loop_param_reps =+        zipWith (\j rep -> (ResTag j, rep)) [next ..] lifted_loop_reps+  (inps_local <> loop_param_inputs,) <$> insertRepsM loop_param_reps env_local++-- transform a for-loop with a variant iteration count into a while-loop+transformForToWhile ::+  FlattenOps ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  StmAux () ->+  [(FParam SOACS, SubExp)] ->+  VName ->+  IntType ->+  SubExp ->+  Body SOACS ->+  FlattenM DistEnv+transformForToWhile ops segments env inps res aux merge i it n body = do+  let old_loop_params = map fst merge+  -- Fresh names used only in the synthetic rewritten body.+  cond_param_v <- newVName "for_cond"+  cond0_v <- newVName "for_cond0"+  cond_next_v <- newVName "for_cond_next"+  i_next_v <- newVName "for_i_next"+  loop_old_out_vs <- replicateM (length merge) $ newVName "for_out"+  i_out_v <- newVName "for_i_out"+  cond_out_v <- newVName "for_cond_out"++  let zero = intConst it 0+      one = intConst it 1+      aux_no_certs = aux {stmAuxCerts = mempty}++      cond0_stm =+        Let+          (Pat [PatElem cond0_v (Prim Bool)])+          aux_no_certs+          (BasicOp $ CmpOp (CmpSlt it) zero n)++      -- Extend the loop parameters with iteration variable and condition variable+      i_param = Param mempty i (Prim (IntType it))+      cond_param = Param mempty cond_param_v (Prim Bool)++      Body loop_body_dec loop_body_stms loop_body_res = body++      i_next_stm =+        Let+          (Pat [PatElem i_next_v (Prim (IntType it))])+          aux_no_certs+          -- OverflowWrap or OverflowUndef?+          (BasicOp $ BinOp (Add it OverflowUndef) (Var i) one)++      cond_next_stm =+        Let+          (Pat [PatElem cond_next_v (Prim Bool)])+          aux_no_certs+          (BasicOp $ CmpOp (CmpSlt it) (Var i_next_v) n)++      loop_new_body =+        Body+          loop_body_dec+          (loop_body_stms <> oneStm i_next_stm <> oneStm cond_next_stm)+          ( [ SubExpRes mempty (Var cond_next_v),+              SubExpRes mempty (Var i_next_v)+            ]+              <> loop_body_res+          )++      merge' =+        [ (cond_param, Var cond0_v),+          (i_param, zero)+        ]+          <> merge++      loop_out_tys = [Prim Bool, Prim (IntType it)] ++ map paramType old_loop_params++      loop_pat =+        Pat $+          zipWith+            PatElem+            ([cond_out_v, i_out_v] ++ loop_old_out_vs)+            loop_out_tys++      while_stm =+        Let+          loop_pat+          aux+          (Loop merge' (WhileLoop (paramName cond_param)) loop_new_body)++      synthetic_body =+        Body+          ()+          (oneStm cond0_stm <> oneStm while_stm)+          (map (SubExpRes mempty . Var) loop_old_out_vs)++  reps <- distributeAndFlattenBody ops segments "for_variant_res" env inps res synthetic_body+  insertRepsM (zip (map distResTag res) reps) env++flattenLoop ::+  FlattenOps ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  (Pat Type, StmAux ()) ->+  ([(Param DeclType, SubExp)], LoopForm, Body SOACS) ->+  FlattenM DistEnv+flattenLoop ops segments env inps res (_pat, aux) (merge, ForLoop i it n, body) = do+  if isVariant inps n+    then transformForToWhile ops segments env inps res aux merge i it n body+    else do+      let old_loop_params = map fst merge+          loopParamNames = S.fromList $ map paramName old_loop_params++      num_segments <- letSubExp "num_segments" =<< toExp (segmentCount segments)+      (lifted_loop_params, lifted_loop_reps, lifted_init) <-+        unzip3 <$> mapM (liftLoopParam (flattenSegLevel ops) segments num_segments inps env loopParamNames) merge++      let lifted_loop_params' = concat lifted_loop_params+          lifted_init' = concat lifted_init++      let i_param = Param mempty i (Prim (IntType it))+          build_scope = scopeOfFParams lifted_loop_params' <> scopeOfLParams [i_param]++      (loop_new_inputs, loop_env_local) <-+        localScope build_scope $+          loopBodyInputs env inps old_loop_params lifted_loop_reps++      -- When the loop parameters and all inputs used by the body are regular,+      -- the interchange of the map-nest and the loop corresponds to a perfectly+      -- ordinary Screma inside the loop. We then transform that Screma as if+      -- that was what the program looked like in the first place, which in+      -- particular means it is subject to multi-versioning. Otherwise we+      -- distribute the loop body statement by statement.+      let body_free = freeIn body+          free_inps =+            [ (v, inp)+            | (v, inp) <- inps,+              v `nameIn` body_free,+              not $ v `S.member` loopParamNames+            ]+          regularInput (_, inp) =+            not (any (variantDim inps loopParamNames) (arrayDims (distInputType inp)))+              && case inp of+                DistInputFree {} -> True+                DistInput rt _ -> case resVar rt env of+                  Regular {} -> True+                  Irregular {} -> False+          regularRep Regular {} = True+          regularRep Irregular {} = False+          simpleParam p = case declTypeOf p of+            Prim {} -> True+            Array {} -> True+            _ -> False+          -- The interchanged Screma is transformed as a top-level+          -- statement, so this is only possible when we are not+          -- generating in-block code.+          at_host_level = case flattenSegLevel ops of+            SegThreadInBlock {} -> False+            _ -> True+          interchangeable =+            at_host_level+              -- Parameters with variant dimensions are lifted to an+              -- irregular representation, so this also rejects those.+              && all regularRep lifted_loop_reps+              && all isRegularDistResult res+              && all simpleParam old_loop_params+              && all regularInput free_inps++      loop_body_gpu <-+        if interchangeable+          then+            localScope build_scope $+              interchangedLoopBody+                ops+                num_segments+                segments+                env+                (zip old_loop_params lifted_loop_params')+                free_inps+                aux+                body+          else+            distributedLoopBody+              ops+              segments+              num_segments+              build_scope+              loop_new_inputs+              loop_env_local+              res+              body++      let loop_exp_gpu =+            Loop+              (zip lifted_loop_params' lifted_init')+              (ForLoop i it n)+              loop_body_gpu++      -- We must copy the result because otherwise we increase the degree of+      -- aliasing. In a loop, the result aliases the input, because it might run+      -- for zero iterations, but in the original program the result was+      -- produced by 'map', which has no aliases.+      loop_out_vs <-+        mapM (letExp "loop_res_out_copy" <=< eCopy . eVar)+          <=< certifying (distCerts inps aux env)+          $ letTupExp "loop_res_out" loop_exp_gpu++      let out_reps = resultToResRepsByDistResult res loop_out_vs+      insertRepsM (zip (map distResTag res) out_reps) env+--+flattenLoop ops segments env inps res (_pat, aux) (merge, WhileLoop cond, body) = do+  -- TODO: Consider updating the active segment so we don't go over w every+  -- time.+  --+  -- inside the body we should compute the indices for which the condition is+  -- true and for which it is false, and then distribute the body based on that.+  -- We can then merge the results of the two branches by writing them back to a+  -- blank space like we do for the branches of a match.+  --+  -- This is probably not worth it: it is faster only for uniform loop+  -- parameters, but the common flattened-while case carries irregular state,+  -- which cannot be narrowed - an irregular result's offsets are global, so+  -- evicting a finished segment needs the final sizes of all segments, which+  -- are not known until the loop ends.++  let old_loop_params = map fst merge+      loopParamNames = S.fromList $ map paramName old_loop_params+  w <- letSubExp "num_segments" =<< toExp (segmentCount segments)+  (lifted_loop_params, lifted_loop_reps, lifted_init) <-+    unzip3 <$> mapM (liftLoopParam lvl segments w inps env loopParamNames) merge++  let lifted_loop_params' = concat lifted_loop_params+      lifted_init' = concat lifted_init+      loop_params_scope = scopeOfFParams lifted_loop_params'+  (loop_new_inputs, loop_env_local) <-+    localScope loop_params_scope $+      loopBodyInputs env inps old_loop_params lifted_loop_reps++  -- find cond_lifted_param in old_lifted_loop_params to get the lifted_loop_reps+  let maybe_cond = lookup cond (zip (map paramName old_loop_params) (zip lifted_loop_reps lifted_init))+  scope <- askScope+  case maybe_cond of+    -- infinite loop+    Nothing -> do+      loop_body_gpu <-+        distributedLoopBody ops segments w loop_params_scope loop_new_inputs loop_env_local res body+      let loop_exp_gpu = Loop (zip lifted_loop_params' lifted_init') (WhileLoop cond) loop_body_gpu+      loop_out_vs <- certifying (distCerts inps aux env) $ letTupExp "loop_res_out" loop_exp_gpu+      let out_reps = resultToResRepsByDistResult res loop_out_vs+      insertRepsM (zip (map distResTag res) out_reps) env+    Just (cond_lifted_rep, cond_init) -> do+      let [cond_init_se] = cond_init++      -- Compute initial any_active+      cond_init_arr_v <- letExp "cond_init_arr" $ BasicOp $ SubExp cond_init_se+      let cond_lifted_param = case cond_lifted_rep of+            Regular v -> v+            Irregular {} -> error "WhileLoop condition cannot be irregular"++      cond_init_arr_t <- lookupType cond_init_arr_v+      cond_init_flat <-+        letExp "cond_init_flat" . BasicOp $+          Reshape cond_init_arr_v $+            reshapeAll (arrayShape cond_init_arr_t) (Shape [w])++      or_lam <- binOpLambda LogOr Bool+      map_lam <- mkIdentityLambda [Prim Bool]+      ~[any_active_init_v] <-+        genNonSegRed lvl "any_active_init" [w] (Reduce Commutative or_lam [constant False]) mempty map_lam [cond_init_flat]+      let any_active_init = Var any_active_init_v++      any_active_param <- newParam "any_active" (Prim Bool)+      let build_scope = loop_params_scope <> scopeOfFParams [any_active_param]+      -- ‌build body+      loop_body_gpu <-+        buildBody_ . localScope build_scope $ do+          -- (num_data, active_inds) <- genFilter cond_lifted_param+          equiv_classes <- letExp "equiv_classes" <=< segMap lvl (MkSolo w) $ \(MkSolo i) -> do+            let seg_is = unflattenIndex (segmentDims segments) (pe64 i)+            c <- letSubExp "c" =<< eIndex cond_lifted_param (map toExp seg_is)+            cls <-+              letSubExp "cls"+                =<< eIf+                  (eSubExp c)+                  (eBody [toExp $ intConst Int64 1])+                  (eBody [toExp $ intConst Int64 0])+            pure [subExpRes cls]+          n_cases <- letExp "n_cases" <=< toExp $ intConst Int64 2+          (partition_sizes, partition_offs, partition_inds) <- doPartition lvl n_cases equiv_classes+          inds_t <- lookupType partition_inds++          let getInds nm k = do+                sz <-+                  letSubExp (nm <> "_sz")+                    =<< eIndex partition_sizes [toExp $ intConst Int64 k]+                off <-+                  letSubExp (nm <> "_off")+                    =<< eIndex partition_offs [toExp $ intConst Int64 k]+                inds <-+                  letExp (nm <> "_inds") $+                    BasicOp $+                      Index partition_inds $+                        fullSlice inds_t [DimSlice off sz (intConst Int64 1)]+                pure (sz, inds)++          (_, inactive_inds) <- getInds "inactive" 0+          (active_size, active_inds) <- getInds "active" 1++          inactive_reps <- forM old_loop_params $ \p -> do+            (_, _, rep) <- splitInput lvl segments loop_env_local loop_new_inputs inactive_inds mempty (paramName p)+            pure rep++          let free_in_body =+                filter+                  (isVariant loop_new_inputs . Var)+                  (namesToList $ freeIn body)+          free_sizes <-+            foldMap freeIn <$> mapM (lookupInputType loop_new_inputs) free_in_body+          let free_variant_sizes = filter (isVariant loop_new_inputs . Var) (namesToList free_sizes)+              free_size_vars = nubOrd (free_variant_sizes <> free_in_body)+          (ts, vs, reps) <- unzip3 <$> mapM (splitInput lvl segments loop_env_local loop_new_inputs active_inds mempty) free_size_vars+          let subset_inputs = do+                (v, t, i) <- zip3 vs ts [0 ..]+                pure (v, DistInput (ResTag i) t)+              env_subset = DistEnv $ M.fromList $ zip (map ResTag [0 ..]) reps+          let subset_segments = [active_size]+          let (subset_inputs', subset_dstms) =+                distributeBodyWith ops scope subset_segments subset_inputs body+          env_subset' <- foldM (flattenDistStm ops subset_segments) env_subset subset_dstms+          active_reps <-+            zipWithM+              (liftDistResultRep lvl subset_segments subset_inputs' env_subset')+              res+              (bodyResult body)++          let mergeOneLifted t rep0 rep1+                | isAcc t = do+                    let (Regular acc_res) = rep1+                    pure [SubExpRes mempty (Var acc_res)]+                | otherwise =+                    case (rep0, rep1) of+                      (Regular x0, Regular x1) -> do+                        let initial_shape = Shape [w] <> arrayShape t+                        let final_shape = segmentsShape segments <> arrayShape t+                        let pt = elemType t+                        space <- letExp "blank" =<< eBlank (Array pt initial_shape NoUniqueness)++                        out <-+                          foldM+                            (scatterRegular lvl)+                            space+                            [(inactive_inds, x0), (active_inds, x1)]++                        out_type <- arrayShape <$> lookupType out+                        out_reshaped <-+                          letExp "out_reshaped" . BasicOp $+                            Reshape out $+                              reshapeAll out_type final_shape++                        pure [SubExpRes mempty (Var out_reshaped)]+                      (Irregular ir0, Irregular ir1) -> do+                        segsSpace <-+                          letExp "blank_segs"+                            =<< eBlank (Array int64 (Shape [w]) NoUniqueness)++                        segs <-+                          foldM+                            (scatterRegular lvl)+                            segsSpace+                            [(inactive_inds, irregularS ir0), (active_inds, irregularS ir1)]++                        (_, offsets, num_data) <- exScanAndSum lvl segs++                        let pt = elemType t+                        elemsSpace <-+                          letExp "blank_elems"+                            =<< eBlank (Array pt (Shape [num_data]) NoUniqueness)++                        elems <-+                          foldM+                            (scatterIrregular lvl offsets)+                            elemsSpace+                            [(inactive_inds, ir0), (active_inds, ir1)]++                        flags <- genFlags lvl num_data offsets++                        pure+                          [ SubExpRes mempty num_data,+                            SubExpRes mempty (Var segs),+                            SubExpRes mempty (Var flags),+                            SubExpRes mempty (Var offsets),+                            SubExpRes mempty (Var elems)+                          ]+                      _ -> error "mergeOneLifted: mismatched reps"++          merged_results <-+            concat+              <$> zipWithM+                (\p (r0, r1) -> mergeOneLifted (declTypeOf p) r0 r1)+                old_loop_params+                (zip inactive_reps active_reps)++          -- we have one extra iteration but it is better than extra reduction in the loop body,+          any_active <-+            letSubExp "any_active"+              =<< eIf+                (toExp $ pe64 active_size .==. 0)+                (eBody [eSubExp $ constant False])+                (eBody [eSubExp $ constant True])++          pure $ merged_results ++ [SubExpRes mempty any_active]++      let merge' =+            zip+              (lifted_loop_params' ++ [any_active_param])+              (lifted_init' ++ [any_active_init])+      loop_out_vs <-+        certifying (distCerts inps aux env) $+          letTupExp "loop_res_out" $+            Loop+              merge'+              (WhileLoop (paramName any_active_param))+              loop_body_gpu+      let loop_out_vs' = L.init loop_out_vs+      let out_reps = resultToResRepsByDistResult res loop_out_vs'+      insertRepsM (zip (map distResTag res) out_reps) env+  where+    lvl = flattenSegLevel ops
+ src/Futhark/Pass/Flatten/Match.hs view
@@ -0,0 +1,328 @@+-- | Flattening of 'Match'.+module Futhark.Pass.Flatten.Match+  ( flattenMatch,+  )+where++import Control.Monad+import Data.Containers.ListUtils (nubOrd)+import Data.List qualified as L+import Data.Map qualified as M+import Data.Maybe+import Data.Set qualified as S+import Data.Tuple.Solo+import Futhark.IR.GPU+import Futhark.IR.SOACS+import Futhark.Pass.Flatten.Distribute+import Futhark.Pass.Flatten.General+import Futhark.Tools++-- Given the indices for which a branch is taken and its body,+-- distribute the statements of the body of that branch.+distributeBranch ::+  FlattenOps ->+  Segments ->+  DistEnv ->+  DistInputs ->+  VName ->+  Body SOACS ->+  M.Map VName ResRep ->+  FlattenM (DistInputs, DistEnv, DistStms)+distributeBranch ops segments env inps is body acc_reps = do+  let lvl = flattenSegLevel ops+      free_in_body = filter (isVariant inps . Var) (namesToList $ freeIn body)+  scope <- askScope+  free_sizes <-+    foldMap freeIn <$> mapM (lookupInputType inps) free_in_body+  let free_variant_sizes = filter (isVariant inps . Var) (namesToList free_sizes)+      free_size_vars = nubOrd (free_variant_sizes <> free_in_body)+  (ts, vs, reps) <-+    unzip3 <$> mapM (splitInput lvl segments env inps is acc_reps) free_size_vars+  let inputs = do+        (v, t, i) <- zip3 vs ts [0 ..]+        pure (v, DistInput (ResTag i) t)+  let env' = DistEnv $ M.fromList $ zip (map ResTag [0 ..]) reps+  let (inputs', dstms) = distributeBodyWith ops scope segments inputs body+  pure (inputs', env', dstms)++-- Given a single result from each branch as well the *unlifted*+-- result type, merge the results of all branches into a single result.+mergeResult ::+  SegLevel ->+  Segments ->+  SubExp ->+  [VName] ->+  [ResRep] ->+  DistResult ->+  FlattenM ResRep+mergeResult lvl segments w iss branchesRep dist_res+  -- Regular case+  | isRegularDistResult dist_res = do+      let (DistType _ _ resType) = distResType dist_res+      if isAcc resType+        then do+          xs <- mapM asRegular branchesRep+          pure $ Regular $ last xs+        else do+          let resultType = Array (elemType resType) (Shape [w] <> arrayShape resType) NoUniqueness+          xs <- mapM asRegular branchesRep+          -- Create the blank space for the result+          resultSpace <- letExp "blank_res" =<< eBlank resultType+          -- Write back the values of each branch to the blank space+          result <- foldM (scatterRegular lvl) resultSpace $ zip iss xs+          result_t <- arrayShape <$> lookupType result+          result' <-+            letExp "match_res_reg" . BasicOp $+              Reshape result (reshapeAll result_t (segmentsShape segments <> arrayShape resType))+          pure $ Regular result'+  -- Irregular case+  | DistType _ _ (Array pt _ _) <- distResType dist_res = do+      branchesIrregRep <- mapM asIrregular branchesRep+      let segsType = Array (IntType Int64) (Shape [w]) NoUniqueness+      -- Create a blank space for the 'segs'+      segsSpace <- letExp "blank_segs" =<< eBlank segsType+      -- Write back the segs of each branch to the blank space+      segs <- foldM (scatterRegular lvl) segsSpace $ zip iss (irregularS <$> branchesIrregRep)+      (_, offsets, num_data) <- exScanAndSum lvl segs+      let resultType = Array pt (Shape [num_data]) NoUniqueness+      -- Create the blank space for the result+      resultSpace <- letExp "blank_res" =<< eBlank resultType+      -- Write back the values of each branch to the blank space+      elems <- foldM (scatterIrregular lvl offsets) resultSpace $ zip iss branchesIrregRep+      flags <- genFlags lvl num_data offsets+      pure $+        Irregular $+          IrregularRep+            { irregularS = segs,+              irregularF = flags,+              irregularO = offsets,+              irregularD = elems,+              irregularK = Dense+            }+  | otherwise = error "mergeResult: non-array irregular result"+  where+    asRegular (Regular v) = pure v+    asRegular _ = error "mergeResult: mismatched reps"++    asIrregular (Irregular irreg) = pure irreg+    asIrregular _ = error "mergeResult: mismatched reps"++-- | Flatten a single branch body of a variant 'Match', but guard its execution+-- on the branch actually being taken by some segment. When a branch receives no+-- segments (its partition is empty) we must not run its flattened code: it may+-- call lifted recursive functions, which would recurse forever on an empty+-- batch. An untaken branch's results are never read ('mergeResult' scatters+-- them back through the branch's empty index array), so we just yield blanks.+--+-- Like 'flattenUniformMatch', the branch is lifted to a 'Result' of flat rep+-- components and the reps recovered with 'distResultsToResReps'; here we+-- additionally wrap it in a @branch_size > 0@ 'Match'.+guardBranch ::+  FlattenOps ->+  SubExp ->+  DistEnv ->+  DistInputs ->+  DistStms ->+  [DistResult] ->+  Result ->+  FlattenM [ResRep]+guardBranch ops branch_size env inputs dstms res result = do+  let branch_segments = [branch_size]+  (taken_body, taken_types) <-+    buildBody $ do+      body_res <- liftBodyWithDistResults ops branch_segments inputs env dstms res result+      ts <- mapM (subExpType . resSubExp) body_res+      pure (body_res, ts)+  -- Blanks for the untaken branch have the same types as the taken branch.+  -- Sizes bound inside the branch (the length of irregular data) are not in+  -- scope here, so we zero them; the 'Match' then makes them existential.+  untaken_body <- buildBody_ $ do+    let blank t = t `setArrayShape` Shape (map (const (intConst Int64 0)) (arrayDims t))+    subExpsRes <$> mapM (letSubExp "blank" <=< eBlank . blank) taken_types+  match_e <-+    eIf+      (eCmpOp (CmpSlt Int64) (eSubExp (intConst Int64 0)) (eSubExp branch_size))+      (pure taken_body)+      (pure untaken_body)+  match_res <- letTupExp "guarded_branch" match_e+  rets <- expExtType match_e+  pure $ distResultsToResReps res $ drop (S.size (shapeContext rets)) match_res++flattenVariantMatch ::+  FlattenOps ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  StmAux () ->+  [SubExp] ->+  [Case (Body SOACS)] ->+  Body SOACS ->+  MatchDec ExtType ->+  FlattenM DistEnv+flattenVariantMatch ops segments env inps res _aux scrutinees cases defaultCase _rt = do+  let lvl = flattenSegLevel ops+  w <- letSubExp "w" <=< toExp $ product $ segmentDims segments+  -- We need to partition the indices of the scrutinees by which case they match.+  -- Lift the scrutinees.+  -- If it's a variable, we know it's a scalar and the lifted version will therefore be a regular array.+  lifted_scrutinees <- forM scrutinees $ \scrut -> do+    liftSubExpRegular lvl segments inps env (segmentsShape segments) scrut+  -- Cases for tagging values that match the same branch.+  -- The default case is the 0'th equvalence class.+  let equiv_cases =+        zipWith+          (\(Case pat _) n -> Case pat $ eBody [toExp $ intConst Int64 n])+          cases+          [1 ..]+  let equiv_case_default = eBody [toExp $ intConst Int64 0]+  -- Match the scrutinees againts the branch cases+  equiv_classes <- letExp "equiv_classes" <=< segMap lvl (MkSolo w) $ \(MkSolo i) -> do+    -- unflatten index+    let seg_is = unflattenIndex (segmentDims segments) (pe64 i)+    scruts <- mapM (letSubExp "scruts" <=< flip eIndex (map toExp seg_is)) lifted_scrutinees+    cls <- letSubExp "cls" =<< eMatch scruts equiv_cases equiv_case_default+    pure [subExpRes cls]+  let num_cases = fromIntegral $ length cases + 1+  n_cases <- letExp "n_cases" <=< toExp $ intConst Int64 num_cases+  -- Parition the indices of the scrutinees by their equvalence class such+  -- that (the indices) of the scrutinees belonging to class 0 come first,+  -- then those belonging to class 1 and so on.+  (partition_sizes, partition_offs, partition_inds) <- doPartition lvl n_cases equiv_classes+  inds_t <- lookupType partition_inds+  -- Get the indices of each scrutinee by equivalence class+  branch_info <- forM [0 .. num_cases - 1] $ \i -> do+    num_data <-+      letSubExp ("size" <> nameFromString (show i))+        =<< eIndex partition_sizes [toExp $ intConst Int64 i]+    begin <-+      letSubExp ("idx_begin" <> nameFromString (show i))+        =<< eIndex partition_offs [toExp $ intConst Int64 i]+    inds <-+      letExp ("inds_branch" <> nameFromString (show i)) $+        BasicOp . Index partition_inds $+          fullSlice inds_t [DimSlice begin num_data (intConst Int64 1)]+    pure (num_data, inds)+  let (branch_sizes, inds) = unzip branch_info++  -- Distribute and lift the branch bodies.+  -- We put the default case at the start as it's the 0'th equivalence class+  -- and is therefore the first segment after the partition.+  let branch_bodies = defaultCase : map (\(Case _ body) -> body) cases+  let branch_results = map bodyResult branch_bodies+  -- Accumulator results are threaded from one branch to the next and cannot be+  -- blanked, so we only guard branch execution when no accumulators are+  -- involved. XXX: can we be sure this will never be a problem?+  let hasAcc = any (\dr -> case distResType dr of DistType _ _ t -> isAcc t) res+  -- acc inputs are handled differently, each branch use the result of the previous branch+  (branch_reps, _) <-+    foldM+      ( \(branch_reps_acc, acc_reps) (branch_size, branch_inds, body, result) -> do+          let branch_segments = [branch_size]+          (inputs, env', dstms) <-+            distributeBranch ops segments env inps branch_inds body acc_reps+          reps <-+            if hasAcc+              then do+                env'' <- foldM (flattenDistStm ops branch_segments) env' dstms+                zipWithM (liftDistResultRep lvl branch_segments inputs env'') res result+              else guardBranch ops branch_size env' inputs dstms res result+          let acc_reps' = replaceAccReps acc_reps reps+          pure (branch_reps_acc <> [reps], acc_reps')+      )+      ([], M.empty)+      (L.zip4 branch_sizes inds branch_bodies branch_results)+  -- Merging acc results is done by using the last branch result+  reps <- zipWithM (mergeResult lvl segments w inds) (L.transpose branch_reps) res+  insertRepsM (zip (map distResTag res) reps) env+  where+    findAccCert :: VName -> (VName, DistInput) -> Maybe VName+    findAccCert cert v_inp =+      let (v, inp) = v_inp+       in if isAcc (distInputType inp)+            then case distInputType inp of+              Acc cert' _ _ _ | cert == cert' -> Just v+              _ -> Nothing+            else Nothing++    -- Idealy this should be a singleton+    findAccCerts :: VName -> [VName]+    findAccCerts cert = mapMaybe (findAccCert cert) inps++    replaceAccRep acc_reps (dist_res, rep) =+      let (DistType _ _ t) = distResType dist_res+       in if not $ isAcc t+            then+              acc_reps+            else+              let (Acc cert _ _ _) = t+                  accVars = findAccCerts cert+               in foldl (\m v -> M.insert v rep m) acc_reps accVars+    replaceAccReps acc_reps reps = foldl replaceAccRep acc_reps $ zip res reps++flattenUniformMatch ::+  FlattenOps ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  StmAux () ->+  [SubExp] ->+  [Case (Body SOACS)] ->+  Body SOACS ->+  MatchDec ExtType ->+  FlattenM DistEnv+flattenUniformMatch ops segments env inps res aux scrutinees cases defaultCase rt = do+  scope <- askScope+  new_cases <- forM cases $ \(Case c body) -> do+    let (case_body_inputs, case_dstms) =+          distributeBodyWith ops scope segments inps body++    fmap (Case c) . buildBody_ $+      liftBodyWithDistResults ops segments case_body_inputs env case_dstms res (bodyResult body)+  new_default_body <- do+    let (new_default_body_inputs, new_default_dstms) =+          distributeBodyWith ops scope segments inps defaultCase+    buildBody_ $+      liftBodyWithDistResults ops segments new_default_body_inputs env new_default_dstms res (bodyResult defaultCase)++  -- Maybe it is better to build MatchDec ourselves+  match_e <-+    eMatch'+      scrutinees+      [Case c (pure body) | Case c body <- new_cases]+      (pure new_default_body)+      (matchSort rt)++  match_res <-+    certifying (distCerts inps aux env) $+      letTupExp "match_res" match_e++  rets <- expExtType match_e+  -- get rid of the existential context+  let payload_res = drop (S.size (shapeContext rets)) match_res+  let reps = distResultsToResReps res payload_res+  insertRepsM (zip (map distResTag res) reps) env++-- | Flatten a 'Match'+flattenMatch ::+  FlattenOps ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  StmAux () ->+  [SubExp] ->+  [Case (Body SOACS)] ->+  Body SOACS ->+  MatchDec ExtType ->+  FlattenM DistEnv+flattenMatch ops segments env inps res aux scrutinees cases defaultCase rt =+  -- 'flattenUniformMatch' keeps the scrutinees in a plain GPU 'Match', which is+  -- only well-scoped when they are invariant to the nest. Whenever a scrutinee is+  -- variant we must partition the segments by branch, even if no branch contains+  -- parallelism (this happens e.g. for a variant conditional with an irregular+  -- result, which cannot be sequentialised into a scalar group).+  if any (isVariant inps) scrutinees+    then flattenVariantMatch ops segments env inps res aux scrutinees cases defaultCase rt+    else flattenUniformMatch ops segments env inps res aux scrutinees cases defaultCase rt
+ src/Futhark/Pass/Flatten/Monad.hs view
@@ -0,0 +1,330 @@+{-# LANGUAGE TypeFamilies #-}++-- | General definitions for the flattening transformation.+--+-- Defines not just the core monads that are involved, but also the various+-- representations, except perhaps the ones that are completely local to another+-- module.+module Futhark.Pass.Flatten.Monad+  ( IrregularKind (..),+    IrregularRep (..),+    ResRep (..),+    DistEnv (..),+    FlattenOps (..),++    -- * Flattening monad+    FlattenM,+    FlattenState (..),+    runFlattenM,++    -- * Demands+    BuiltinFn (..),+    LiftMode (..),+    DemandFn (..),+    demandLifted,+    demandBuiltin,++    -- * Insertions+    insertRepM,+    insertRepsM,+    insertIrregularM,+    insertRegulars,++    -- * Various+    resVar,+    inputReps,+    segsAndElems,+  )+where++import Control.Monad+import Control.Monad.State+import Data.Bifunctor (bimap, second)+import Data.Foldable+import Data.Map qualified as M+import Data.Maybe (fromMaybe)+import Data.Set qualified as S+import Futhark.IR.GPU+import Futhark.IR.SOACS (SOACS)+import Futhark.Pass.Flatten.Distribute+import Futhark.Tools+import Prelude hiding (div, rem)++-- | If true, 'sanityCheck' blocks are evaluated.+doSanityCheck :: Bool+doSanityCheck = True++-- | Run a sanity-check that may verify invariants. The idea is that these can+-- be disabled without affecting the correctness of the pass, although there is+-- no constructive guarantee that no important effects take place in here.+sanityCheck :: (Monad m) => m () -> m ()+sanityCheck = when doSanityCheck++-- Note [Representation of Flat Arrays]+--+-- This flattening implementation uses largely the nomenclature and+-- structure described by Cosmin Oancea. In particular, consider an+-- irregular array 'A' where+--+--   - A has 'n' segments (outermost dimension).+--+--   - A has element type 't'.+--+--   - A has a total of 'm' elements (where 'm' is divisible by 'n',+--     and may indeed be 'm').+--+-- Then A is represented by the following arrays:+--+--   - A_D : [m]t; the "data array".+--+--   - A_S : [n]i64; the "shape array" giving the number of scalar elements of each segment.+--+--   - A_F : [m]bool; the "flag array", indicating when an element begins a+--     new segment.+--+--   - A_O : [n]i64; the offset array, indicating for each segment+--     where it starts in the data (and flag) array.+--+--   - A_II1 : [m]i64; the "segment indices"; a mapping from element+--     index to index of the segment it belongs to.+--+--   - A_II2 : [m]i64; the "inner indices"; a mapping from element index+--     to index within its corresponding segment.+--+-- The arrays that are not the data array are collectively called the+-- "structure arrays". All of the structure arrays can be computed+-- from each other, but conceptually they all coexist.+--+-- Note that we only consider the *outer* dimension to be the+-- "segments". Also, 't' may actually be an array itself (although in+-- this case, the shape of 't' must be invariant to all parallel+-- dimensions). The inner structure is preserved through code, not+-- data. (Or in practice, ad-hoc auxiliary arrays produced by code.)+-- In Cosmin's notation, we maintain only the information for the+-- outermost dimension.+--+-- As an example, consider an irregular array+--+--   A = [ [], [ [1,2,3], [4], [], [5,6] ], [ [7], [], [8,9,10] ] ]+--+-- then+--+--   n = 3+--+--   m = 10+--+--   A_D = [1,2,3,4,5,6,7,8,9,10]+--+--   A_S = [0, 6, 4]+--+--   A_F = [T,F,F,F,F,F,T,F,F,F]+--+--   A_O = [0, 0, 6]+--+--   A_II1 = [1,1,1,1,1,1,2,2,2,2]+--+--   A_II2 = [0,0,0,1,3,3,0,2,2,2]++data IrregularKind+  = Dense+  | Replicated+  deriving (Show, Eq)++data IrregularRep = IrregularRep+  { -- | Array of size of each segment, type @[]i64@.+    irregularS :: VName,+    irregularF :: VName,+    irregularO :: VName,+    irregularD :: VName,+    irregularK :: IrregularKind+  }+  deriving (Show)++data ResRep+  = -- | This variable is represented completely straightforwardly- if it is an+    -- array, it is a regular array.+    Regular VName+  | -- | The representation of an irregular array.+    Irregular IrregularRep+  deriving (Show)++newtype DistEnv = DistEnv {distResMap :: M.Map ResTag ResRep}++insertRep :: ResTag -> ResRep -> DistEnv -> DistEnv+insertRep rt rep env = env {distResMap = M.insert rt rep $ distResMap env}++insertRepM :: ResTag -> ResRep -> DistEnv -> FlattenM DistEnv+insertRepM rt rep env = do+  sanityCheck $ do+    case rep of+      Regular _ -> pure ()+      Irregular (IrregularRep shape flags offsets data_ _kind) -> do+        shape_t <- lookupType shape+        flags_t <- lookupType flags+        data_t <- lookupType data_+        offsets_t <- lookupType offsets++        unless (arrayRank flags_t == 1 && elemType flags_t == Bool) $+          error $+            "Invalid flag array type: " <> prettyString flags_t+        unless (arrayRank offsets_t == 1 && elemType offsets_t == int64) $+          error $+            "Invalid offsets array type: " <> prettyString offsets_t+        unless (arrayRank shape_t == 1 && elemType shape_t == int64) $+          error $+            "Invalid shape array type: " <> prettyString shape_t+        when (arrayRank data_t /= 1) $+          error $+            "Invalid data array array: " <> prettyString data_t+  pure $ insertRep rt rep env++insertRepsM :: [(ResTag, ResRep)] -> DistEnv -> FlattenM DistEnv+insertRepsM =+  flip $ foldM (flip $ uncurry insertRepM)++insertReps :: [(ResTag, ResRep)] -> DistEnv -> DistEnv+insertReps = flip $ foldl (flip $ uncurry insertRep)++insertIrregularM :: VName -> VName -> VName -> ResTag -> VName -> IrregularKind -> DistEnv -> FlattenM DistEnv+insertIrregularM shape flags offsets rt data_ kind env = do+  let rep = Irregular $ IrregularRep shape flags offsets data_ kind+  insertRepM rt rep env++insertRegulars :: [ResTag] -> [VName] -> DistEnv -> DistEnv+insertRegulars rts xs =+  insertReps (zip rts $ map Regular xs)++instance Monoid DistEnv where+  mempty = DistEnv mempty++instance Semigroup DistEnv where+  DistEnv x <> DistEnv y = DistEnv (x <> y)++resVar :: ResTag -> DistEnv -> ResRep+resVar rt env = fromMaybe bad $ M.lookup rt $ distResMap env+  where+    bad = error $ "resVar: unknown tag: " ++ show rt++segsAndElems :: DistEnv -> [DistInput] -> (Maybe (VName, VName, VName), [VName])+segsAndElems _ [] = (Nothing, [])+segsAndElems env (DistInputFree v _ : vs) =+  second (v :) $ segsAndElems env vs+segsAndElems env (DistInput rt _ : vs) =+  case resVar rt env of+    Regular v' ->+      second (v' :) $ segsAndElems env vs+    Irregular (IrregularRep segments flags offsets elems k) -> do+      case k of+        Dense -> do+          bimap (mplus $ Just (segments, flags, offsets)) (elems :) $ segsAndElems env vs+        Replicated ->+          second (flags :) $ segsAndElems env vs++-- | Mapping from original variable names to their distributed resreps.+inputReps :: DistInputs -> DistEnv -> M.Map VName (Type, ResRep)+inputReps inputs env = M.fromList $ map (second getRep) inputs+  where+    getRep di = case di of+      DistInput rt t -> (t, resVar rt env)+      DistInputFree v' t -> (t, Regular v')++-- | A representation of the different kinds of builtin functions we can+-- generate. This is used to only generate the ones we actually need for a given+-- program.+data BuiltinFn+  = BuiltinSegIota+  | BuiltinRepIota+  | BuiltinPrefixSum+  | BuiltinPartition+  deriving (Eq, Ord, Show)++data LiftMode+  = UniformLift+  | NonUniformLift+  deriving (Eq, Ord, Show)++-- | Indicate the need for a function to be generated. Instead of immediately+-- generating them ourselves, we collect requirements from multiple flattening+-- operations and satisfy them in their entirety.+data DemandFn+  = -- | We need this function to be lifted.+    DemandLifted Name LiftMode+  | DemandBuiltin BuiltinFn+  deriving (Eq, Ord, Show)++data FlattenState = FlattenState+  { -- In order to generate more stable threshold names, we keep track of the+    -- numbers used for thresholds separately from the ordinary name source.+    stateThresholdCounter :: Int,+    stateNameSource :: VNameSource,+    -- A set of those functions that we have emitted calls to, and which will+    -- need to be generated.+    stateDemandFns :: S.Set DemandFn+  }++newtype FlattenM a = FlattenM (BuilderT GPU (State FlattenState) a)+  deriving+    ( Functor,+      Applicative,+      Monad,+      LocalScope GPU,+      HasScope GPU,+      MonadState FlattenState,+      MonadFreshNames+    )++instance MonadBuilder FlattenM where+  type Rep FlattenM = GPU+  mkExpDecM pat e = FlattenM $ mkExpDecM pat e+  mkBodyM stms res = FlattenM $ mkBodyM stms res+  mkLetNamesM pat e = FlattenM $ mkLetNamesM pat e++  addStms = FlattenM . addStms+  collectStms (FlattenM m) = FlattenM $ collectStms m++instance MonadFreshNames (State FlattenState) where+  getNameSource = gets stateNameSource+  putNameSource src = modify $ \s -> s {stateNameSource = src}++-- | Do not nest these - the counter for thresholds will be wrong.+runFlattenM :: (MonadFreshNames m) => Scope GPU -> FlattenM a -> m (a, S.Set DemandFn)+runFlattenM scope (FlattenM m) = modifyNameSource $ \src ->+  let initial_state =+        FlattenState+          { stateThresholdCounter = 0,+            stateNameSource = src,+            stateDemandFns = mempty+          }+      (x, s) = runState (fst <$> runBuilderT m scope) initial_state+   in ((x, stateDemandFns s), stateNameSource s)++-- | Indicate that we rather need a lifted version of this function.+demandLifted :: Name -> LiftMode -> FlattenM ()+demandLifted fname mode = modify $ \s ->+  s {stateDemandFns = S.insert (DemandLifted fname mode) $ stateDemandFns s}++-- | Demand the presence of this builtin function.+demandBuiltin :: BuiltinFn -> FlattenM ()+demandBuiltin b = modify $ \s ->+  s {stateDemandFns = S.insert (DemandBuiltin b) $ stateDemandFns s}++-- | Functions for tying together disparate modules - this is to avoid mutually+-- recursive modules.+data FlattenOps = FlattenOps+  { flattenSegLevel :: SegLevel,+    -- | How to treat irregularity when distributing the bodies we encounter.+    flattenIrregularity :: DistIrregularity,+    flattenFunHasParallelism :: FunHasParallelism,+    -- | Flatten a 'DistStm' using the given ops, which need not be the ones+    -- this record belongs to - see 'atSegLevel'. Use 'flattenDistStm' to+    -- continue with the current ops.+    flattenDistStmWith :: FlattenOps -> Segments -> DistEnv -> DistStm -> FlattenM DistEnv,+    -- | Flatten a scalar statement at the given seg level, which need not be+    -- the one this record carries. Use 'flattenScalarStm' for the latter.+    flattenScalarStmAt :: SegLevel -> Segments -> DistEnv -> DistInputs -> [DistResult] -> Stm SOACS -> FlattenM DistEnv,+    -- | Transform a statement as if it occurred at the top level, including+    -- multi-versioning of SOACs. Used when a transformation (e.g. loop+    -- interchange) produces a statement that should be treated as if the+    -- program had looked like that all along.+    flattenTopLevelStm :: Stm SOACS -> FlattenM ()+  }
+ src/Futhark/Pass/Flatten/PreProcess.hs view
@@ -0,0 +1,136 @@+{-# LANGUAGE TypeFamilies #-}++-- | Preprocess the program before flattening.  This rewrites SOAC forms+-- that flatten does not want to see directly, while leaving the result in+-- SOACS form so the normal flattening pipeline can continue afterwards.+module Futhark.Pass.Flatten.PreProcess+  ( shouldDissectForm,+    preprocessProg,+    preprocessBody,+    preprocessStms,+    preprocessStm,+    preprocessLambda,+    runSimplifiedBuilder,+  )+where++import Data.Maybe (isNothing)+import Futhark.Builder+import Futhark.IR.SOACS+import Futhark.IR.SOACS.Simplify+import Futhark.Pass+import Futhark.Tools+import Futhark.Transform.FirstOrderTransform qualified as FOT+import Futhark.Transform.ISRWIM (irwim, iswim)++shouldDissectForm :: ScremaForm SOACS -> Bool+shouldDissectForm form =+  isNothing (isMapSOAC form)+    && isNothing (isReduceSOAC form)+    && isNothing (isScanSOAC form)+    && isNothing (isRedomapSOAC form)+    && isNothing (isScanomapSOAC form)+    && isNothing (isMaposcanomapSOAC form)++runSimplifiedBuilder ::+  (MonadFreshNames m) =>+  Scope SOACS ->+  BuilderT SOACS m a ->+  m (Stms SOACS)+runSimplifiedBuilder scope m =+  fst <$> runBuilderT (simplifyStms =<< collectStms_ m) scope++-- | Rewrite a SOAC form that flattening does not handle directly into one it+-- does, recursively preprocessing the result (which may itself contain further+-- SOACs). Returns 'Nothing' for statements that need no rewriting - those are+-- handled structurally by 'preprocessStm'.+rewriteSoacStm ::+  (MonadFreshNames m) =>+  Scope SOACS ->+  Stm SOACS ->+  Maybe (m (Stms SOACS))+rewriteSoacStm scope (Let pat aux (Op soac))+  | "sequential_outer" `inAttrs` stmAuxAttrs aux =+      Just $+        preprocessStms scope =<< runSimplifiedBuilder scope (FOT.transformSOAC pat soac)+rewriteSoacStm scope (Let pat aux (Op (Stream w arrs nes lam))) = Just $ do+  stms <- runSimplifiedBuilder scope (auxing aux $ sequentialStreamWholeArray pat w nes lam arrs)+  preprocessStms scope stms+rewriteSoacStm scope (Let pat aux (Op (Screma w' arrs' form')))+  | Just scans <- isScanSOAC form',+    Scan scan_lam nes <- singleScan scans,+    Just do_iswim <- iswim pat w' scan_lam (zip nes arrs') = Just $ do+      stms <- runSimplifiedBuilder scope $ auxing aux do_iswim+      preprocessStms scope stms+  | Just [Reduce comm red_fun nes] <- isReduceSOAC form',+    let comm'+          | commutativeLambda red_fun = Commutative+          | otherwise = comm,+    Just do_irwim <- irwim pat w' comm' red_fun (zip nes arrs') = Just $ do+      stms <- runSimplifiedBuilder scope $ auxing aux do_irwim+      preprocessStms scope stms+  | shouldDissectForm form' = Just $ do+      stms <- runSimplifiedBuilder scope (auxing aux $ dissectScrema pat w' form' arrs')+      preprocessStms scope stms+rewriteSoacStm _ _ = Nothing++preprocessStm ::+  (MonadFreshNames m) =>+  Scope SOACS ->+  Stm SOACS ->+  m (Stms SOACS)+preprocessStm _ stm+  | "sequential" `inAttrs` stmAuxAttrs (stmAux stm) = pure $ oneStm stm+preprocessStm scope stm+  | Just rewritten <- rewriteSoacStm scope stm = rewritten+preprocessStm scope (Let pat aux (Loop merge form body)) = do+  let scope' = scopeOfFParams (map fst merge) <> scopeOfLoopForm form <> scope+  body' <- preprocessBody scope' body+  pure $ oneStm $ Let pat aux $ Loop merge form body'+preprocessStm scope (Let pat aux (Match ses cases defbody dec)) = do+  cases' <- mapM (traverse (preprocessBody scope)) cases+  defbody' <- preprocessBody scope defbody+  pure $ oneStm $ Let pat aux $ Match ses cases' defbody' dec+preprocessStm scope (Let pat aux (WithAcc inputs lam)) = do+  lam' <- preprocessLambda scope lam+  pure $ oneStm $ Let pat aux $ WithAcc inputs lam'+preprocessStm _ stm = pure $ oneStm stm++preprocessStms ::+  (MonadFreshNames m) =>+  Scope SOACS ->+  Stms SOACS ->+  m (Stms SOACS)+preprocessStms scope stms = mconcat <$> mapM (preprocessStm scope') (stmsToList stms)+  where+    scope' = scopeOf stms <> scope++preprocessBody ::+  (MonadFreshNames m) =>+  Scope SOACS ->+  Body SOACS ->+  m (Body SOACS)+preprocessBody scope body = do+  stms <- preprocessStms scope $ bodyStms body+  pure $ body {bodyStms = stms}++preprocessLambda ::+  (MonadFreshNames m) =>+  Scope SOACS ->+  Lambda SOACS ->+  m (Lambda SOACS)+preprocessLambda scope lam = do+  body <- preprocessBody (scopeOfLParams (lambdaParams lam) <> scope) $ lambdaBody lam+  let lam' = lam {lambdaBody = body}+  fst <$> runBuilderT (simplifyLambda lam') scope++preprocessFun :: Stms SOACS -> FunDef SOACS -> PassM (FunDef SOACS)+preprocessFun consts fd = do+  body <- preprocessBody (scopeOf consts <> scopeOf fd) $ funDefBody fd+  pure $ fd {funDefBody = body}++preprocessProg :: Prog SOACS -> PassM (Prog SOACS)+preprocessProg =+  intraproceduralTransformationWithConsts+    (preprocessStms mempty)+    preprocessFun
+ src/Futhark/Pass/Flatten/SOAC.hs view
@@ -0,0 +1,1920 @@+{-# LANGUAGE LambdaCase #-}++-- | Flattening rules for SOACs.+module Futhark.Pass.Flatten.SOAC+  ( flattenScrema,+    flattenHist,+    flattenFlatMap,+    flattenFlatMapNested,+  )+where++import Control.Monad+import Data.Containers.ListUtils (nubOrd)+import Data.Foldable+import Data.Functor.Identity (runIdentity)+import Data.Map qualified as M+import Data.Maybe (fromMaybe, mapMaybe)+import Data.Tuple.Solo+import Futhark.IR.GPU+import Futhark.IR.SOACS+import Futhark.MonadFreshNames+import Futhark.Pass.Flatten.Distribute+import Futhark.Pass.Flatten.General+import Futhark.Pass.Flatten.Incremental+import Futhark.Pass.Flatten.Intrablock qualified as Intrablock+import Futhark.Pass.Flatten.PreProcess+import Futhark.Tools+import Futhark.Transform.FirstOrderTransform qualified as FOT+import Futhark.Transform.Rename+import Futhark.Transform.Substitute+import Futhark.Transform.ToGPU (soacsLambdaToGPU)+import Futhark.Util (mapAccumLM)+import Futhark.Util.IntegralExp+import Prelude hiding (div, quot, rem)++-- | How the results of a nested map are represented, determined by whether the+-- map width is uniform (invariant to the enclosing nest) or nonuniform; see+-- 'transformInnerMap'.+data InnerMapMode+  = -- | Uniform width: the results are regular arrays that keep the full+    -- multi-dimensional shape (the enclosing segments followed by the map+    -- width).+    MultiDim+  | -- | Nonuniform width: the results are irregular, flattened into a single+    -- segment dimension.+    SingleDim++freeWithTypeDeps :: DistInputs -> Names -> FlattenM [VName]+freeWithTypeDeps inps free = do+  let free_names = namesToList free+  free_sizes <- foldMap freeIn <$> mapM (lookupInputType inps) free_names+  pure $ nubOrd $ namesToList free_sizes <> free_names++-- Reduction or scan operators may not have any free variables that are variant+-- to the nest (that is, are inputs to the distributed operation), and must+-- operate on primitive types. This is because we would be unable to express+-- them as SegScan/SegReds. Fixing this would require modifications to the+-- SegOp representation, but it is likely not worth it, as such operators are+-- extremely rare - and we can just fall back on sequentialising the SOAC and+-- flattening the resulting loop.+suitableOperator :: DistEnv -> DistInputs -> Lambda SOACS -> [SubExp] -> Bool+suitableOperator _env inps lam _nes =+  allNames notVariant (freeIn lam)+    && all primType (lambdaReturnType lam)+  where+    notVariant = not . isVariant inps . Var++suitableUniformOperator :: DistInputs -> Lambda SOACS -> Bool+suitableUniformOperator inps lam =+  allNames (not . isVariant inps . Var) (freeIn lam)++regularToReplicatedIrregularRep ::+  SegLevel ->+  Segments ->+  VName ->+  VName ->+  FlattenM IrregularRep+regularToReplicatedIrregularRep lvl segments ws_data v' = do+  ws_prod <- arraySize 0 <$> lookupType ws_data+  arr_t <- lookupType v'+  segment_size <-+    letSubExp "reg_seg_size" <=< toExp . product . map pe64 $+      drop (segmentsRank segments) (arrayDims arr_t)+  num_elems <-+    letSubExp "reg_num_elems" <=< toExp $ product $ map pe64 $ arrayDims arr_t+  arr_D <-+    letExp "reg_D" . BasicOp $+      Reshape v' (reshapeAll (arrayShape arr_t) (Shape [num_elems]))+  arr_F <- letExp "reg_F" <=< segMap lvl (MkSolo num_elems) $ \(MkSolo i) -> do+    flag <- letSubExp "flag" <=< toExp $ (pe64 i `rem` pe64 segment_size) .==. 0+    pure [subExpRes flag]++  arr_S <-+    letExp "reg_segments" . BasicOp $+      Replicate (Shape [ws_prod]) segment_size+  arr_O <- letExp "reg_O" <=< segMap lvl (MkSolo ws_prod) $ \(MkSolo i) -> do+    segment <- letSubExp "segment" =<< eIndex ws_data [eSubExp i]+    offset <- letSubExp "offset" <=< toExp $ pe64 segment * pe64 segment_size+    pure [subExpRes offset]+  let rep' =+        IrregularRep+          { irregularS = arr_S,+            irregularF = arr_F,+            irregularO = arr_O,+            irregularD = arr_D,+            irregularK = Replicated+          }+  pure rep'++-- Replicates inner dimension for inputs.+onMapFreeVar ::+  SegLevel ->+  Segments ->+  DistEnv ->+  DistInputs ->+  VName ->+  (VName, VName, VName) ->+  VName ->+  Maybe (FlattenM (VName, MapArray IrregularRep))+onMapFreeVar lvl segments env inps _ws (_ws_F, _ws_O, ws_data) v = do+  v_inp <- lookup v inps+  pure $ do+    ws_prod <- arraySize 0 <$> lookupType ws_data+    fmap (v,) $ case v_inp of+      DistInputFree v' t ->+        --  I'm not totally sure if this will be better than previous approach+        (`MapOther` t) <$> regularToReplicatedIrregularRep lvl segments ws_data v'+      DistInput rt t -> case resVar rt env of+        Irregular rep -> do+          ~[new_S, offsets] <- letTupExp (baseName v <> "_rep_free_irreg")+            <=< segMap lvl (MkSolo ws_prod)+            $ \(MkSolo i) -> do+              segment <- letSubExp "segment" =<< eIndex ws_data [eSubExp i]+              s <- letSubExp "s" =<< eIndex (irregularS rep) [eSubExp segment]+              o <- letSubExp "o" =<< eIndex (irregularO rep) [eSubExp segment]+              pure $ subExpsRes [s, o]+          let rep' =+                IrregularRep+                  { irregularS = new_S,+                    irregularF = irregularF rep,+                    irregularO = offsets,+                    irregularD = irregularD rep,+                    irregularK = Replicated+                  }+          pure $ MapOther rep' t+        Regular vs ->+          (`MapOther` t) <$> regularToReplicatedIrregularRep lvl segments ws_data vs++onMapFreeVarMultiDim ::+  SegLevel ->+  Segments ->+  SubExp ->+  DistEnv ->+  DistInputs ->+  VName ->+  Maybe (FlattenM (VName, MapArray IrregularRep))+onMapFreeVarMultiDim lvl segments w env inps v = do+  v_inp <- lookup v inps+  pure $ fmap (v,) $ case v_inp of+    DistInputFree v' t -> do+      v_rep <- replicateForDims segments (Shape [w]) v'+      pure $ MapArray v_rep t+    DistInput rt t -> case resVar rt env of+      Regular v' -> do+        v_rep <- replicateForDims segments (Shape [w]) v'+        pure $ MapArray v_rep t+      Irregular rep -> do+        -- Can replicate as well+        old_nseg <- arraySize 0 <$> lookupType (irregularS rep)+        new_nseg <- letSubExp "new_nseg" <=< toExp $ pe64 old_nseg * pe64 w+        ~[new_S, offsets] <- letTupExp (baseName v <> "_rep_free_irreg")+          <=< segMap lvl (MkSolo new_nseg)+          $ \(MkSolo i) -> do+            old_seg <- letSubExp "old_seg" <=< toExp $ pe64 i `quot` pe64 w+            s <- letSubExp "s" =<< eIndex (irregularS rep) [eSubExp old_seg]+            o <- letSubExp "o" =<< eIndex (irregularO rep) [eSubExp old_seg]+            pure $ subExpsRes [s, o]+        let rep' =+              IrregularRep+                { irregularS = new_S,+                  irregularF = irregularF rep,+                  irregularO = offsets,+                  irregularD = irregularD rep,+                  irregularK = Replicated+                }+        pure $ MapOther rep' t++onMapInputArr ::+  SegLevel ->+  Segments ->+  DistEnv ->+  DistInputs ->+  VName ->+  VName ->+  VName ->+  Param Type ->+  VName ->+  FlattenM (MapArray IrregularRep)+onMapInputArr lvl segments env inps ws ws_O ws_data p arr = do+  ws_prod <- arraySize 0 <$> lookupType ws_data+  case lookup arr inps of+    Just v_inp ->+      case v_inp of+        DistInputFree vs t -> do+          let inner_shape = arrayShape $ paramType p+          vs_t <- lookupType vs+          v <-+            if isAcc vs_t+              then pure vs+              else+                letExp (baseName vs <> "_flat") . BasicOp . Reshape vs $+                  reshapeAll (arrayShape vs_t) (Shape [ws_prod] <> inner_shape)+          pure $ MapArray v (rowType t)+        DistInput rt t ->+          case resVar rt env of+            Irregular rep -> do+              onMapIrregularInputArr lvl SingleDim segments ws ws_O ws_data p arr rep ws_prod+            Regular vs -> do+              let inner_shape = arrayShape $ paramType p+              vs_t <- lookupType vs+              if isAcc vs_t+                then pure $ MapArray vs t+                else do+                  v <-+                    letExp (baseName arr <> "_reg_flat") . BasicOp . Reshape vs $+                      reshapeAll (arrayShape vs_t) (Shape [ws_prod] <> inner_shape)+                  pure $ MapArray v (stripArray 1 vs_t)+    Nothing -> do+      arr_row_t <- rowType <$> lookupType arr+      arr_rep <-+        letExp (baseName arr <> "_inp_rep") . BasicOp $+          Replicate (segmentsShape segments) (Var arr)+      arr_rep_t <- lookupType arr_rep+      v <-+        letExp (baseName arr <> "_inp_rep_flat") . BasicOp . Reshape arr_rep $+          reshapeAll (arrayShape arr_rep_t) (Shape [ws_prod] <> arrayShape arr_row_t)+      pure $ MapArray v arr_row_t++mapArraysToInputs ::+  [VName] ->+  [MapArray IrregularRep] ->+  FlattenM (DistEnv, DistInputs)+mapArraysToInputs param_names arrs = do+  ((_, env), inputs) <-+    mapAccumLM onInput (0, mempty) $ zip param_names arrs+  pure (env, inputs)+  where+    onInput (tag, env) (p, MapArray arr t) =+      pure ((tag, env), (p, DistInputFree arr t))+    onInput (tag, env) (p, MapOther rep t) = do+      let rt = ResTag tag+      env' <- insertRepM rt (Irregular rep) env+      pure+        ( (tag + 1, env'),+          (p, DistInput rt t)+        )++transformUniformRedomap ::+  SegLevel ->+  Segments ->+  DistEnv ->+  DistInputs ->+  SubExp ->+  [VName] ->+  [Reduce SOACS] ->+  Lambda SOACS ->+  FlattenM [VName]+transformUniformRedomap lvl [] _env _inps w arrs reds map_lam = do+  -- Top-level (no enclosing segments): the arrays and any free variables are+  -- ordinary top-level values in scope, so this is an ordinary non-segmented+  -- reduce over the map width. We emit it as such ('genNonSegRed'); a segmented+  -- reduce over a single implicit segment would be equivalent, but downstream+  -- passes (e.g. migration, coalescing) handle the non-segmented form better.+  let sing_red = singleReduce reds+  (red_lam, nes', shape) <- determineReduceOp (redLambda sing_red) (redNeutral sing_red)+  let comm+        | commutativeLambda red_lam = Commutative+        | otherwise = redComm sing_red+      sing_red_gpu = Reduce comm (soacsLambdaToGPU red_lam) nes'+  genNonSegRed lvl "topLevelSegRed" [w] sing_red_gpu shape (soacsLambdaToGPU map_lam) arrs+transformUniformRedomap lvl segments env inps w arrs reds map_lam = do+  let free = freeIn map_lam+      new_segment = segments <> pure w+      shape = mempty+  reds_gpu <- forM reds $ \red -> do+    nes <- mapM (readNeutral segments env inps) (redNeutral red)+    let red_lam = redLambda red+        comm+          | commutativeLambda red_lam = Commutative+          | otherwise = redComm red+    pure $ Reduce comm (soacsLambdaToGPU red_lam) nes+  free_and_sizes <- freeWithTypeDeps inps free+  (free_replicated, replicated) <-+    fmap unzip . sequence $+      mapMaybe+        (onMapFreeVarMultiDim lvl segments w env inps)+        free_and_sizes+  arrs' <-+    zipWithM+      ( \p arr ->+          liftSubExpRegular+            lvl+            segments+            inps+            env+            (segmentsShape new_segment <> arrayShape (paramType p))+            (Var arr)+      )+      (lambdaParams map_lam)+      arrs++  (free_env, free_inputs) <- mapArraysToInputs free_replicated replicated+  let readFree is = readInputs new_segment free_env is free_inputs+  genUniformSegRed lvl "uniformSegRed" new_segment reds_gpu shape (soacsLambdaToGPU map_lam) arrs' readFree++doUniformSegMaposcanomap ::+  SegLevel ->+  [Scan SOACS] ->+  [VName] ->+  Lambda SOACS ->+  Lambda SOACS ->+  Segments ->+  Segments ->+  DistInputs ->+  DistEnv ->+  ([SubExp] -> FlattenM ()) ->+  FlattenM [VName]+doUniformSegMaposcanomap lvl scans arrs post_lam map_lam old_segments new_segment inps env readFree = do+  let scan = singleScan scans+  nes <- mapM (readNeutral old_segments env inps) (scanNeutral scan)+  (scan_lam, nes', shape) <- determineReduceOp (scanLambda scan) nes+  genUniformSegScanomapWithPost+    lvl+    new_segment+    "uniformmaposcanomap"+    (soacsLambdaToGPU scan_lam)+    shape+    nes'+    (soacsLambdaToGPU post_lam)+    (soacsLambdaToGPU map_lam)+    arrs+    readFree++transformUniformMaposcanomap ::+  SegLevel ->+  Segments ->+  DistEnv ->+  DistInputs ->+  SubExp ->+  [VName] ->+  [Scan SOACS] ->+  Lambda SOACS ->+  Lambda SOACS ->+  FlattenM [VName]+transformUniformMaposcanomap lvl segments env inps w arrs scans post_lam map_lam = do+  let free = freeIn map_lam <> freeIn post_lam+      new_segment = segments <> pure w+  free_and_sizes <- freeWithTypeDeps inps free+  (free_replicated, replicated) <-+    fmap unzip . sequence $+      mapMaybe+        (onMapFreeVarMultiDim lvl segments w env inps)+        free_and_sizes+  arrs' <-+    zipWithM+      ( \p arr ->+          liftSubExpRegular+            lvl+            segments+            inps+            env+            (segmentsShape new_segment <> arrayShape (paramType p))+            (Var arr)+      )+      (lambdaParams map_lam)+      arrs+  (free_env, free_inputs) <- mapArraysToInputs free_replicated replicated+  let readFree is = readInputs new_segment free_env is free_inputs+  doUniformSegMaposcanomap lvl scans arrs' post_lam map_lam segments new_segment inps env readFree++doSegMaposcanomap ::+  SegLevel ->+  [Scan SOACS] ->+  VName ->+  [VName] ->+  Lambda SOACS ->+  Lambda SOACS ->+  Segments ->+  DistInputs ->+  DistEnv ->+  ([SubExp] -> FlattenM ()) ->+  FlattenM [VName]+doSegMaposcanomap lvl scans flags elems post_lam map_lam segments inps env readFree = do+  let scan = singleScan scans+  let nes = scanNeutral scan+  nes' <- mapM (readNeutral segments env inps) nes+  genSegScanomapWithPost+    lvl+    "maposcanomap"+    (soacsLambdaToGPU $ scanLambda scan)+    nes'+    flags+    (soacsLambdaToGPU post_lam)+    (soacsLambdaToGPU map_lam)+    elems+    readFree++-- Hacky fix to get result representations in the same order as the pattern+resRepsInPatOrder :: Pat Type -> [(VName, ResRep)] -> [ResRep]+resRepsInPatOrder pat reps =+  let rep_map = M.fromList reps+      lookupRes v =+        case M.lookup v rep_map of+          Just rep -> rep+          Nothing ->+            error $+              "resRepsInPatOrder: missing result for "+                ++ prettyString v+   in map lookupRes (patNames pat)++segOpInputRep ::+  SegLevel ->+  Segments ->+  DistEnv ->+  DistInputs ->+  VName ->+  FlattenM ResRep+segOpInputRep lvl segments env inps arr =+  case lookup arr inps of+    Just (DistInput rt _) ->+      pure $ resVar rt env+    Just (DistInputFree arr' _) ->+      pure $ Regular arr'+    Nothing ->+      Irregular <$> getIrregRep lvl segments env inps arr++-- Basically we need to make our arrays ready for our segscan/segred.+-- Regular arrays are flattened only across the outer segment dimensions and+-- the SOAC width; any row shape expected by the consumer is preserved.+-- we need to check the dense/replicated status of the input.+-- if all of scan inputs are replicated we are fine.+-- otherwise, we need to make the replicated inputs dense.+-- for regulars we can just use the segment descriptor and this should be also the same descriptor for dense irregulars.+prepareSegOpInputs ::+  SegLevel ->+  Segments ->+  DistEnv ->+  DistInputs ->+  SubExp ->+  [ResRep] ->+  [VName] ->+  Bool ->+  FlattenM (VName, VName, VName, [VName], IrregularKind)+prepareSegOpInputs lvl segments env inps w reps names hasNoFreeVariant+  | all isRegular reps = do+      ws <- dataArr lvl segments env inps w+      (ws_F, ws_O, ws_data) <- doRepIota lvl ws+      m <- arraySize 0 <$> lookupType ws_data+      names' <- mapM (flattenRegularRep m) reps+      pure (ws_F, ws_O, ws, names', Dense)+  | all isReplicatedIrregular reps && hasNoFreeVariant = do+      -- We use the descriptor of the first rep for all inputs, which assumes+      -- that all the replicated inputs have the same offsets into their+      -- respective data arrays. This holds because same-width views produced+      -- by onMapFreeVar inherit the offsets of their underlying arrays, and+      -- those are compact per-segment arrays of the SOAC width (slices are+      -- materialised before they can become inputs here).+      let Irregular rep0 = head reps+      pure (irregularF rep0, irregularO rep0, irregularS rep0, map getData reps, Replicated)+  | otherwise = do+      -- The segment descriptor must count SOAC *elements* per segment, but an+      -- 'IrregularRep' stores primitive data whose structure arrays count+      -- _scalars_ - @c@ per element for a non-scalar element type @b=[c]t@. We+      -- take a dense input rep as the descriptor and convert its scalar-unit+      -- structure to element units: sizes and offsets divided by that input's+      -- @c@, flags subsampled with stride @c@ (all no-ops when @c=1@). Each+      -- input's flat data is reshaped into @[m]b@ rows.+      row_types <- mapM (fmap rowType . lookupInputType inps) names+      (desc, desc_c) <- descriptor $ zip reps row_types+      ws_S <- scaleSizesDown lvl desc_c (irregularS desc)+      ws_O <- scaleSizesDown lvl desc_c (irregularO desc)+      ws_F <- subsampleFlags lvl desc_c (irregularF desc)+      m <- arraySize 0 <$> lookupType ws_F+      names' <- sequence $ zipWith3 (toRows m) reps names row_types+      pure (ws_F, ws_O, ws_S, names', Dense)+  where+    isRegular (Regular _) = True+    isRegular _ = False++    isReplicatedIrregular (Irregular rep) = irregularK rep == Replicated+    isReplicatedIrregular _ = False++    flattenRegularRep m (Regular v) =+      flattenRegularToRows segments m v+    flattenRegularRep _ _ =+      error "prepareSegOpInputs: impossible irregular regular input"+    getData (Irregular rep) = irregularD rep+    getData _ = error "prepareSegOpInputs: impossible"++    -- A dense input rep to use as the descriptor, and its element type's inner+    -- size @c@. Any dense irregular input serves, since all inputs share the+    -- SOAC's segmentation.+    descriptor rs =+      case [(rep, t) | (Irregular rep, t) <- rs, irregularK rep == Dense] of+        (rep, t) : _ -> pure (rep, innerSize t)+        [] ->+          case [(rep, t) | (Irregular rep, t) <- rs] of+            (rep, t) : _ -> do+              rep' <- ensureDenseIrregular lvl "segop_desc" rep+              pure (rep', innerSize t)+            [] -> error "prepareSegOpInputs: impossible"+    innerSize t = product $ map pe64 $ arrayDims t++    -- Reshape an input's data into @m@ rows of the (possibly non-scalar) element+    -- type. For an irregular input the flat scalar data is grouped into rows;+    -- for a regular input the enclosing dimensions are collapsed.+    toRows m (Regular v') _ _ =+      flattenRegularToRows segments m v'+    toRows m (Irregular ir) v row_t = do+      d <- irregularD <$> ensureDenseIrregular lvl (baseName v <> "_dense") ir+      d_t <- lookupType d+      letExp (baseName v <> "_rows") . BasicOp $+        Reshape d $+          reshapeAll (arrayShape d_t) (Shape [m] <> arrayShape row_t)++flattenRegularToRows :: Segments -> SubExp -> VName -> FlattenM VName+flattenRegularToRows segments m v = do+  v_t <- lookupType v+  if isAcc v_t+    then pure v+    else do+      when (arrayRank v_t < segmentsRank segments + 1) $+        error "prepareSegOpInputs: regular input rank too small"+      let row_shape = arrayShape $ stripArray (segmentsRank segments + 1) v_t+      letExp (baseName v <> "_flat") . BasicOp $+        Reshape v $+          reshapeAll (arrayShape v_t) (Shape [m] <> row_shape)++-- | Construct a body and immediately rename it.+renamedBody :: FlattenM [VName] -> FlattenM (Body GPU)+renamedBody = renameBody <=< buildBody_ . fmap varsRes++regularRepVars :: [ResRep] -> [VName]+regularRepVars =+  map onRep+  where+    onRep (Regular v) = v+    onRep Irregular {} = error "regularRepVars: expected regular result"++versionScanRed ::+  FlattenOps ->+  Name ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  StmAux () ->+  SubExp ->+  Body SOACS ->+  FlattenM [VName] ->+  FlattenM DistEnv+versionScanRed ops desc segments env inps res aux w factored_body outer_only = do+  let result_ts =+        [ t `arrayOfShape` segmentsShape segments+        | DistResult _ (DistType _ _ t) _ <- res+        ]+  outer_body <- renamedBody outer_only+  full_body <- case segments of+    -- Top-level (no enclosing segments): flatten the factored body's statements+    -- as ordinary top-level statements. Unlike distributing them over segments,+    -- this copes with array-valued operators and nested SOACs whose temporaries+    -- would otherwise escape the segmented machinery's scope.+    [] ->+      renameBody <=< buildBody_ $ do+        mapM_ (flattenTopLevelStm ops) $ bodyStms factored_body+        pure $ bodyResult factored_body+    _ ->+      renamedBody $ regularRepVars <$> distributeAndFlattenBody ops segments "versionScanRed_full_body" env inps res factored_body++  match_res <-+    certifying (distCerts inps aux env) $+      scanRedAlternatives+        desc+        result_ts+        (stmAuxAttrs aux)+        (isParallelFunInside (flattenFunHasParallelism ops) factored_body)+        (allowVersioning (flattenSegLevel ops))+        (segments <> pure w)+        full_body+        outer_body+  pure $ insertRegulars (map distResTag res) match_res env++insertSegOpMapResults ::+  Segments ->+  VName ->+  VName ->+  VName ->+  IrregularKind ->+  [(DistResult, VName)] ->+  DistEnv ->+  FlattenM DistEnv+insertSegOpMapResults segments segs flags offsets kind bnds env0 =+  foldM insert env0 bnds+  where+    insert env (dist_res, v)+      | isRegularDistResult dist_res = do+          let DistType _ _ t = distResType dist_res+          if isAcc t+            then pure $ insertRegulars [distResTag dist_res] [v] env+            else do+              let expected_shape = segmentsShape segments <> arrayShape t+              v_t <- lookupType v+              v' <-+                letExp (baseName v <> "_reshaped") . BasicOp $+                  Reshape v $+                    reshapeAll (arrayShape v_t) expected_shape+              pure $ insertRegulars [distResTag dist_res] [v'] env+      | otherwise =+          insertIrregularM segs flags offsets (distResTag dist_res) v kind env++distResCerts :: DistEnv -> [DistInput] -> Certs+distResCerts env = Certs . map f+  where+    f (DistInputFree v _) = v+    f (DistInput rt _) = case resVar rt env of+      Regular v -> v+      Irregular r -> irregularD r++reshapeAndBind :: VName -> VName -> Shape -> FlattenM ()+reshapeAndBind v src shape = do+  v_copy <- letExp (baseName v) . BasicOp $ SubExp $ Var src+  v_copy_shape <- arrayShape <$> lookupType v_copy+  letBindNames [v] $ BasicOp $ Reshape v_copy $ reshapeAll v_copy_shape shape++mapResultRep :: SegLevel -> InnerMapMode -> (VName, VName, VName) -> VName -> FlattenM ResRep+mapResultRep _ MultiDim _ v = pure $ Regular v+mapResultRep lvl SingleDim (ws, ws_F, ws_O) v =+  -- Forcing the irregular rep to be 1D because in some places that is my assumption+  -- and also this will make the metadata consistent.+  Irregular+    <$> flattenIrregularRep+      lvl+      IrregularRep+        { irregularS = ws,+          irregularF = ws_F,+          irregularO = ws_O,+          irregularD = v,+          irregularK = Dense+        }++transformDistributed ::+  FlattenOps ->+  M.Map ResTag IrregularRep ->+  Segments ->+  Distributed ->+  FlattenM ()+transformDistributed ops irregs segments dist = do+  let Distributed dstms (DistResults resmap reps) = dist+  env <- foldM (flattenDistStm ops segments) env_initial dstms+  forM_ (M.toList resmap) $ \(rt, binds) ->+    forM_ binds $ \(cs_inps, v, v_t) ->+      certifying (distResCerts env cs_inps) $+        case resVar rt env of+          Regular v' -> letBindNames [v] $ BasicOp $ SubExp $ Var v'+          Irregular irreg -> do+            -- It might have an irregular representation, but we know+            -- that it is actually regular because it is a result.+            irreg' <- ensureDenseIrregular (flattenSegLevel ops) (baseName v <> "_dist_res") irreg+            reshapeAndBind v (irregularD irreg') (segmentsShape segments <> arrayShape v_t)+  forM_ reps $ \(v, r) ->+    case r of+      Left se ->+        letBindNames [v] $ BasicOp $ Replicate (segmentsShape segments) se+      Right (DistInputFree arr _) ->+        letBindNames [v] $ BasicOp $ SubExp $ Var arr+      -- This can happen. ask Troels+      Right (DistInput rt t) ->+        case resVar rt env of+          Regular v' -> letBindNames [v] $ BasicOp $ SubExp $ Var v'+          Irregular irreg -> do+            irreg' <- ensureDenseIrregular (flattenSegLevel ops) (baseName v <> "_dist_rep") irreg+            reshapeAndBind v (irregularD irreg') (segmentsShape segments <> arrayShape t)+  where+    env_initial = DistEnv {distResMap = M.map Irregular irregs}++onMapIrregularInputArr ::+  SegLevel ->+  InnerMapMode ->+  Segments ->+  VName ->+  VName ->+  VName ->+  Param Type ->+  VName ->+  IrregularRep ->+  SubExp ->+  FlattenM (MapArray IrregularRep)+onMapIrregularInputArr lvl mode new_segments ws ws_O ws_data p arr rep ws_prod = do+  -- new_segments already has the new w inside, unlike other functions+  rep_t <- lookupType $ irregularD rep+  let p_t = paramType p+  when (arrayRank rep_t > 1) $+    error "onMapIrregularInputArr: irregularD is not 1D"+  if null (arrayDims p_t)+    then do+      -- Assuming irregularD is 1D, size(irregularD rep) == ws_prod should hold and this should be fine.+      let old_shape = arrayShape rep_t+          new_shape =+            case mode of+              SingleDim -> Shape [ws_prod]+              MultiDim -> segmentsShape new_segments+      case irregularK rep of+        Dense -> do+          v_reshaped <- letExp (baseName (paramName p) <> "_reshaped") $ BasicOp $ Reshape (irregularD rep) $ reshapeAll old_shape new_shape+          pure $ MapArray v_reshaped p_t+        Replicated -> do+          new_flat <-+            letExp (baseName arr <> "_flat_expand")+              <=< segMap lvl (MkSolo ws_prod)+              $ \(MkSolo i) -> do+                j <- letSubExp "j" =<< eIndex ws_data [eSubExp i]+                data_off <- letSubExp "data_off" =<< eIndex (irregularO rep) [eSubExp j]+                seg_start <- letSubExp "seg_start" =<< eIndex ws_O [eSubExp j]+                local_pos <- letSubExp "local_pos" <=< toExp $ pe64 i - pe64 seg_start+                flat_idx <- letSubExp "flat_idx" <=< toExp $ pe64 data_off + pe64 local_pos+                fmap (subExpsRes . pure) $ letSubExp "elem" =<< eIndex (irregularD rep) [eSubExp flat_idx]+          v_reshaped <- letExp (baseName (paramName p) <> "_reshaped") $ BasicOp $ Reshape new_flat $ reshapeAll old_shape new_shape+          pure $ MapArray v_reshaped p_t+    else do+      -- We need to split multi-dimensional irregular segments into per-row+      -- segments. We compute the per-row size by dividing each segment's total+      -- size by its number of rows. The division is exact: within a single+      -- segment the array is an ordinary rectangular value, so all rows have+      -- the same size - irregularity exists only across segments. The+      -- alternative would be to read the row size from the sizes in the+      -- parameter type, but those are per-segment distributed inputs, and we do+      -- not have the environment at hand here to look them up.+      num_segments <- arraySize 0 <$> lookupType ws+      -- per_row_size[s] = irregularS[s] / ws[s]+      per_row_size <-+        letExp (baseName (paramName p) <> "_per_row_size")+          <=< segMap lvl (MkSolo num_segments)+          $ \(MkSolo s) -> do+            total_s <- letSubExp "total_s" =<< eIndex (irregularS rep) [eSubExp s]+            num_rows_s <- letSubExp "num_rows_s" =<< eIndex ws [eSubExp s]+            row_size <-+              letSubExp "row_size"+                =<< eIf+                  (toExp $ pe64 num_rows_s .==. 0)+                  (eBody [toExp $ intConst Int64 0])+                  (eBody [toExp $ pe64 total_s `div` pe64 num_rows_s])+            pure $ subExpsRes [row_size]+      new_S <-+        letExp (baseName (paramName p) <> "_new_S")+          <=< segMap lvl (MkSolo ws_prod)+          $ \(MkSolo i) -> do+            seg_i <- letSubExp "seg_i" =<< eIndex ws_data [eSubExp i]+            sz <- letSubExp "sz" =<< eIndex per_row_size [eSubExp seg_i]+            pure $ subExpsRes [sz]+      rep' <- case irregularK rep of+        Dense -> do+          (_, new_O, m) <- exScanAndSum lvl new_S+          new_F <- genFlags lvl m new_O+          pure $+            IrregularRep+              { irregularD = irregularD rep,+                irregularF = new_F,+                irregularS = new_S,+                irregularO = new_O,+                irregularK = Dense+              }+        Replicated -> do+          new_O <-+            letExp (baseName (paramName p) <> "_new_O")+              <=< segMap lvl (MkSolo ws_prod)+              $ \(MkSolo i) -> do+                seg_i <- letSubExp "seg_i" =<< eIndex ws_data [eSubExp i]+                row_size <- letSubExp "row_size" =<< eIndex per_row_size [eSubExp seg_i]+                seg_row_start <- letSubExp "seg_row_start" =<< eIndex ws_O [eSubExp seg_i]+                row_in_seg <- letSubExp "row_in_seg" <=< toExp $ pe64 i - pe64 seg_row_start+                base_off <- letSubExp "base_off" =<< eIndex (irregularO rep) [eSubExp seg_i]+                off <- letSubExp "off" <=< toExp $ pe64 base_off + pe64 row_in_seg * pe64 row_size+                pure $ subExpsRes [off]+          m <- arraySize 0 <$> lookupType (irregularD rep)+          -- we will have mutliple write but it is the same value so it should be fine.+          new_F <- genFlags lvl m new_O+          pure $+            IrregularRep+              { irregularD = irregularD rep,+                irregularF = new_F,+                irregularS = new_S,+                irregularO = new_O,+                irregularK = Replicated+              }+      pure $ MapOther rep' p_t++onMapInputArrMultiDim ::+  SegLevel ->+  Segments ->+  SubExp ->+  DistEnv ->+  DistInputs ->+  VName ->+  VName ->+  VName ->+  Param Type ->+  VName ->+  FlattenM (MapArray IrregularRep)+onMapInputArrMultiDim lvl old_segments w env inps ws ws_O ws_data p arr = do+  case lookup arr inps of+    Just v_inp ->+      case v_inp of+        DistInputFree vs t -> pure $ MapArray vs (rowType t)+        DistInput rt t -> case resVar rt env of+          Irregular rep -> do+            ws_prod <- arraySize 0 <$> lookupType ws_data+            onMapIrregularInputArr lvl MultiDim (old_segments <> pure w) ws ws_O ws_data p arr rep ws_prod+          Regular vs -> do+            vs_t <- lookupType vs+            if isAcc vs_t+              then pure $ MapArray vs (rowType t)+              else do+                -- let's be cautious and make sure it has the correct shape+                let expected_shape = segmentsShape old_segments <> arrayShape t+                if arrayShape vs_t == expected_shape+                  then pure $ MapArray vs t+                  else do+                    v <-+                      letExp (baseName arr <> "_reg_reshape") . BasicOp . Reshape vs $+                        reshapeAll (arrayShape vs_t) expected_shape+                    pure $ MapArray v (rowType t)+    Nothing -> do+      arr_row_t <- rowType <$> lookupType arr+      arr_rep <-+        letExp (baseName arr <> "_inp_rep") . BasicOp $+          Replicate (segmentsShape old_segments) (Var arr)+      pure $ MapArray arr_rep arr_row_t++flattenMapForInBlock ::+  FlattenOps ->+  Pat Type ->+  SubExp ->+  [VName] ->+  Lambda SOACS ->+  FlattenM ()+flattenMapForInBlock ops pat w arrs map_lam = do+  scope <- askScope+  lam <- preprocessLambda (castScope scope) map_lam+  let arrs' = zipWith MapArray arrs $ map paramType (lambdaParams lam)+      (distributed, _) =+        distributeMapWith ops' scope pat [w] arrs' lam+  transformDistributed ops' mempty [w] distributed+  where+    ops' = atSegLevel inBlockSegLevel ops++resultMapMode :: InnerMapMode -> DistInputs -> Type -> InnerMapMode+resultMapMode SingleDim _ _ = SingleDim+resultMapMode MultiDim new_inps v_t+  | any (isVariant new_inps) (arrayDims v_t) = SingleDim+  | otherwise = MultiDim++irregularMapResult ::+  SegLevel ->+  InnerMapMode ->+  (VName, VName, VName) ->+  Segments ->+  IrregularRep ->+  VName ->+  Type ->+  DistInputs ->+  FlattenM ResRep+irregularMapResult lvl mode (ws, ws_F, ws_O) segments irreg v v_t new_inps =+  do+    irreg_dense <- ensureDenseIrregular lvl (baseName v <> "_map_result") irreg+    if any (isVariant new_inps) (arrayShape v_t)+      then do+        old_segment <- arraySize 0 <$> lookupType ws+        -- The size of each flattened outer segment is the sum of its rows'+        -- sizes. Because irreg_dense is dense (compact offsets), we get this in+        -- O(1) per segment as last_offset + last_size - start, avoiding a+        -- segmented reduction over the row sizes. The guard handles empty outer+        -- segments, which have no last row to read (and would index out of+        -- bounds).+        new_shape <- letExp (baseName v <> "_outer_shape") <=< segMap lvl (MkSolo old_segment) $ \(MkSolo is) -> do+          outer_ind <- letSubExp "outer_ind" =<< eIndex ws_O [eSubExp is]+          outer_ws_i <- letSubExp "outer_ws" =<< eIndex ws [eSubExp is]+          sz <-+            letSubExp "sz"+              =<< eIf+                (toExp $ pe64 outer_ws_i .==. 0)+                (eBody [toExp $ intConst Int64 0])+                ( do+                    last_row <- letSubExp "last_row" <=< toExp $ pe64 outer_ind + pe64 outer_ws_i - 1+                    start <- letSubExp "start" =<< eIndex (irregularO irreg_dense) [eSubExp outer_ind]+                    last_offset <- letSubExp "last_offset" =<< eIndex (irregularO irreg_dense) [eSubExp last_row]+                    last_size <- letSubExp "last_size" =<< eIndex (irregularS irreg_dense) [eSubExp last_row]+                    eBody [toExp $ pe64 last_offset - pe64 start + pe64 last_size]+                )+          pure [subExpRes sz]+        (new_ws_F, new_ws_O, _) <- doRepIota lvl new_shape+        letBindNames [v] $ BasicOp $ Replicate mempty $ Var $ irregularD irreg_dense+        mapResultRep lvl SingleDim (new_shape, new_ws_F, new_ws_O) v+      else do+        reshapeAndBind v (irregularD irreg_dense) (segmentsShape segments <> arrayShape v_t)+        mapResultRep lvl mode (ws, ws_F, ws_O) v++transformDistributedInnerMap ::+  FlattenOps ->+  InnerMapMode ->+  (VName, VName, VName) ->+  M.Map ResTag IrregularRep ->+  Segments ->+  Distributed ->+  FlattenM [(VName, ResRep)]+transformDistributedInnerMap ops mode (ws_F, ws_O, ws) irregs segments dist = do+  let Distributed dstms (DistResults resmap reps) = dist+  -- A name bound inside the distributed body is variant whether or not another+  -- statement uses it, so a result sized by such a name is irregular. Only+  -- counting the names that are used would leave the size existentially bound+  -- by a 'FlatMap' lambda (returned, but used by nothing) looking invariant,+  -- and the result would be given a type mentioning a name that is not in+  -- scope.+  let asInput (DistResult tag (DistType _ _ t) v) = (v, DistInput tag t)+      new_inps =+        concatMap distStmInputs dstms+          <> map asInput (concatMap distStmResult dstms)+  env <- foldM (flattenDistStm ops segments) env_initial dstms+  resmap_res <- fmap concat $ forM (M.toList resmap) $ \(rt, binds) ->+    forM binds $ \(cs_inps, v, v_t) ->+      certifying (distResCerts env cs_inps) $+        case (resultMapMode mode new_inps v_t, resVar rt env) of+          (MultiDim, Regular v') ->+            if isAcc v_t+              then do+                letBindNames [v] $ BasicOp $ SubExp $ Var v'+                pure (v, Regular v)+              else do+                reshapeAndBind v v' (segmentsShape segments <> arrayShape v_t)+                pure (v, Regular v)+          (SingleDim, Regular v') ->+            if isAcc v_t+              then do+                letBindNames [v] $ BasicOp $ SubExp $ Var v'+                pure (v, Regular v)+              else do+                letBindNames [v] $ BasicOp $ SubExp $ Var v'+                rep <- mapResultRep lvl SingleDim (ws, ws_F, ws_O) v+                pure (v, rep)+          (result_mode, Irregular irreg) -> do+            rep <- irregularMapResult lvl result_mode (ws, ws_F, ws_O) segments irreg v v_t new_inps+            pure (v, rep)+  reps_res <- forM reps $ \(v, r) -> do+    case r of+      Left se -> do+        letBindNames [v] $ BasicOp $ Replicate (segmentsShape segments) se+        -- the se is not part of input so this should be fine+        rep <- mapResultRep lvl mode (ws, ws_F, ws_O) v+        pure (v, rep)+      Right (DistInputFree arr t) -> do+        letBindNames [v] $ BasicOp $ SubExp $ Var arr+        if isAcc t+          then pure (v, Regular v)+          else do+            rep <- mapResultRep lvl (resultMapMode mode new_inps t) (ws, ws_F, ws_O) v+            pure (v, rep)+      Right (DistInput rt t) ->+        let result_mode = resultMapMode mode new_inps t+         in case resVar rt env of+              Regular v' -> do+                letBindNames [v] $ BasicOp $ SubExp $ Var v'+                if isAcc t+                  then pure (v, Regular v)+                  else do+                    rep <- mapResultRep lvl result_mode (ws, ws_F, ws_O) v+                    pure (v, rep)+              Irregular irreg -> do+                rep <- irregularMapResult lvl result_mode (ws, ws_F, ws_O) segments irreg v t new_inps+                pure (v, rep)+  pure $ resmap_res <> reps_res+  where+    env_initial = DistEnv {distResMap = M.map Irregular irregs}+    lvl = flattenSegLevel ops++distributeAndTransformInnerMap ::+  FlattenOps ->+  InnerMapMode ->+  (VName, VName, VName) ->+  Segments ->+  DistInputs ->+  Pat Type ->+  [MapArray IrregularRep] ->+  (VName -> Maybe (FlattenM (VName, MapArray IrregularRep))) ->+  Lambda SOACS ->+  FlattenM [ResRep]+distributeAndTransformInnerMap ops mode ws_triple new_segment inps pat arrs' onFreeVar map_lam = do+  -- Skip the return type: the results are described by 'pat', so variables+  -- occurring only there never need replicating into the body - and they may+  -- not even be in scope here. See Note [Ill-formed inner-map lambda].+  let free = freeIn $ map_lam {lambdaReturnType = [] :: [Type]}+  free_and_sizes <- freeWithTypeDeps inps free+  (free_replicated, replicated) <-+    fmap unzip . sequence $+      mapMaybe+        onFreeVar+        free_and_sizes+  free_ps <-+    zipWithM+      newParam+      (map ((<> "_free") . baseName) free_replicated) -- this should free_replicated?+      (map mapArrayRowType replicated)+  scope <- askScope+  let substs = M.fromList $ zip free_replicated $ map paramName free_ps+      map_lam' =+        substituteNames+          substs+          ( map_lam+              { lambdaParams = free_ps <> lambdaParams map_lam+              }+          )+      (distributed, arrmap) =+        distributeMapWith ops scope pat new_segment (replicated <> arrs') map_lam'+  -- order the result representations in the same order as the pattern+  resRepsInPatOrder pat+    <$> transformDistributedInnerMap ops mode ws_triple arrmap new_segment distributed++-- | Flatten a map nested in a map-nest (nonempty enclosing 'Segments'). The map+-- width is either uniform (invariant to the nest) or nonuniform, giving the+-- 'InnerMapMode': a uniform width produces a regular, multi-dimensional result+-- ('MultiDim'), while a nonuniform width is irregular and flattened into a+-- single segment dimension ('SingleDim'). The mode selects how inputs and free+-- variables are read and how the results are represented, and is threaded+-- through the rest of the flattening.+transformInnerMap ::+  FlattenOps ->+  Segments ->+  DistEnv ->+  DistInputs ->+  Pat Type ->+  SubExp ->+  [VName] ->+  Lambda SOACS ->+  FlattenM [ResRep]+transformInnerMap ops segments env inps pat w arrs map_lam = do+  outer_scope <- askScope+  let mode+        | isVariant inps w = SingleDim+        | otherwise = MultiDim+      -- In the uniform 'MultiDim' case - regular inputs and all result+      -- dimensions invariant - the flags/offsets/elements bookkeeping produced+      -- by 'doRepIota' is never consulted, so we do not emit it. This is not+      -- just an efficiency concern: when generating in-block code, the+      -- bookkeeping contains SegOps whose dimensions are bound inside the kernel+      -- body, which would make 'noNonuniformPar' reject the enclosing intrablock+      -- version.+      invariantDim Constant {} = True+      invariantDim (Var v) = v `M.member` outer_scope+      regularInput arr = case lookup arr inps of+        Just (DistInput rt _)+          | Irregular {} <- resVar rt env -> False+        _ -> True+      uniform =+        all (all invariantDim . arrayDims) (patTypes pat)+          && all regularInput arrs+  (ws, ws_F, ws_O, ws_data) <-+    case mode of+      MultiDim+        | uniform ->+            -- XXX: this depends on laziness to explode only on usage. It might+            -- be better to handle this path more explicitly.+            pure (bad "ws", bad "ws_F", bad "ws_O", bad "ws_data")+      _ -> do+        ws <- dataArr lvl segments env inps w+        (ws_F, ws_O, ws_data) <- doRepIota lvl ws+        pure (ws, ws_F, ws_O, ws_data)+  (arrs', new_segment, onFreeVar) <-+    case mode of+      MultiDim -> do+        arrs' <-+          zipWithM+            (onMapInputArrMultiDim lvl segments w env inps ws ws_O ws_data)+            (lambdaParams map_lam)+            arrs+        pure (arrs', segments <> pure w, onMapFreeVarMultiDim lvl segments w env inps)+      SingleDim -> do+        arrs' <-+          zipWithM+            (onMapInputArr lvl segments env inps ws ws_O ws_data)+            (lambdaParams map_lam)+            arrs+        new_segment <- arraySize 0 <$> lookupType ws_data+        pure (arrs', [new_segment], onMapFreeVar lvl segments env inps ws (ws_F, ws_O, ws_data))+  distributeAndTransformInnerMap ops mode (ws_F, ws_O, ws) new_segment inps pat arrs' onFreeVar map_lam+  where+    lvl = flattenSegLevel ops+    bad what =+      error $ "transformInnerMap: " <> what <> " demanded in uniform case"++-- | Flatten a map over the given enclosing 'Segments'. With no enclosing+-- segments this is a top-level map, whose inputs are ordinary regular values;+-- otherwise it is a map nested in a map-nest, whose inputs are the+-- per-enclosing-segment values.+transformMap ::+  FlattenOps ->+  -- | Incremental-flattening attributes of the enclosing statement, propagated+  -- onto the (preprocessed) body in the top-level case; see+  -- 'transformTopLevelMap'.+  Attrs ->+  Segments ->+  DistEnv ->+  DistInputs ->+  Pat Type ->+  SubExp ->+  [VName] ->+  Lambda SOACS ->+  FlattenM [ResRep]+transformMap ops attrs [] _env _inps pat w arrs map_lam = do+  -- Top-level map (no enclosing segments). Preprocess the body and then+  -- propagate the enclosing attributes onto it, so they influence how the body+  -- is versioned (e.g. only_inner reaching a Screma produced by interchanging a+  -- 'sequential_outer' loop). Order matters: preprocessing may rewrite a body+  -- statement, so propagating first would lose the attributes on the rewritten+  -- form. XXX: this is arguably a bug in preprocessing.+  scope <- castScope <$> askScope :: FlattenM (Scope SOACS)+  lam <-+    fmap (propagateVersioningAttrs attrs) . renameLambda+      =<< preprocessLambda scope map_lam+  transformTopLevelMap ops pat w arrs lam+transformMap ops _attrs segments env inps pat w arrs map_lam = do+  gpu_scope <- askScope+  let pp_scope = castScope $ scopeOfDistInputs inps <> gpu_scope+  lam <- preprocessLambda pp_scope map_lam+  transformInnerMap ops segments env inps pat w arrs lam++-- | Fully flatten a map that has no enclosing segments (a top-level map). This+-- is the empty-'Segments' special case of 'transformMap': the mapped arrays are+-- ordinary regular top-level values, so we distribute the map directly over its+-- own width and flatten the resulting body, rather than reconstructing+-- per-enclosing-segment inputs. The results are necessarily regular.+transformTopLevelMap ::+  FlattenOps ->+  Pat Type ->+  SubExp ->+  [VName] ->+  Lambda SOACS ->+  FlattenM [ResRep]+transformTopLevelMap ops pat w arrs lam = do+  scope <- castScope <$> askScope :: FlattenM (Scope SOACS)+  let arrs' = zipWith MapArray arrs $ map paramType (lambdaParams lam)+      (distributed, _) =+        distributeMapWith ops scope pat [w] arrs' lam+  transformDistributed ops mempty [w] distributed+  pure $ map Regular $ patNames pat++runMapLambdaBody ::+  Segments ->+  DistEnv ->+  DistInputs ->+  SubExp ->+  [VName] ->+  Lambda SOACS ->+  Pat Type ->+  [DistResult] ->+  FlattenM [VName]+runMapLambdaBody [] _env _inps w arrs map_lam _pat _ress = do+  -- Top level (no enclosing segments): the mapped arrays are indexed directly+  -- and free variables are already in scope, so there is no per-segment input+  -- reconstruction to do - just run the (sequentialised) body under a segmap+  -- over the map width.+  map_lam' <- renameLambda $ soacsLambdaToGPU map_lam+  vs <- letTupExp "outer_map" <=< renameExp <=< segMap defaultSegLevel [w] $ \is -> do+    let gtid = case toList is of+          [i] -> i+          _ -> error "runMapLambdaBody: expected single index"+    forM_ (zip (lambdaParams map_lam') arrs) $ \(p, arr) ->+      letBindNames [paramName p]+        =<< case paramType p of+          Acc {} -> eSubExp $ Var arr+          _ -> eIndex arr [eSubExp gtid]+    bodyBind $ lambdaBody map_lam'+  forM vs $ \v ->+    letExp (baseName v <> "_copy") $ BasicOp $ Replicate mempty (Var v)+runMapLambdaBody segments env inps w arrs map_lam _pat _ress = do+  map_lam' <- renameLambda $ soacsLambdaToGPU map_lam+  ws <- dataArr defaultSegLevel segments env inps w+  (_ws_F, ws_O, ws_data) <- doRepIota defaultSegLevel ws+  arrs' <-+    zipWithM+      (onMapInputArrMultiDim defaultSegLevel segments w env inps ws ws_O ws_data)+      (lambdaParams map_lam')+      arrs++  free_and_sizes <- freeWithTypeDeps inps (freeIn map_lam')+  (param_env, param_inputs) <-+    mapArraysToInputs (map paramName (lambdaParams map_lam')) arrs'+  let new_segments = segments <> pure w+      free_inputs =+        [ (v, inp)+        | v <- free_and_sizes,+          Just inp <- [lookup v inps]+        ]++  vs <- letTupExp "outer_map" <=< renameExp <=< segMap defaultSegLevel new_segments $ \is -> do+    let full_is = toList is+        outer_is = take (segmentsRank segments) full_is++    readInputs segments env outer_is free_inputs+    readInputs new_segments param_env full_is param_inputs++    bodyBind $ lambdaBody map_lam'+  forM vs $ \v -> do+    letExp (baseName v <> "_copy") $+      BasicOp $+        Replicate mempty (Var v)++versionedUniformMap ::+  FlattenOps ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  Pat Type ->+  StmAux () ->+  SubExp ->+  [VName] ->+  Lambda SOACS ->+  FlattenM DistEnv+versionedUniformMap ops segments env inps ress pat aux w arrs map_lam = do+  let only_intra = onlyExploitIntra (stmAuxAttrs aux)+      may_intra = worthIntrablock map_lam && mayExploitIntra (stmAuxAttrs aux)++  intra' <-+    if only_intra || may_intra+      then Intrablock.intrablockParallelise (flattenMapForInBlock ops) segments env inps ress pat aux w arrs map_lam+      else pure Nothing++  let fullFlatten =+        regularRepVars <$> transformMap (atSegLevel defaultSegLevel ops) (stmAuxAttrs aux) segments env inps pat w arrs map_lam++      outerOnly =+        runMapLambdaBody segments env inps w arrs map_lam pat ress++  full_body <- renamedBody fullFlatten+  outer_body <- renamedBody outerOnly++  let result_ts =+        [ t `arrayOfShape` segmentsShape segments+        | DistResult _ (DistType _ _ t) _ <- ress+        ]++  match_res <-+    certifying (distCerts inps aux env) $+      mapAlternatives+        "match_res"+        result_ts+        (stmAuxAttrs aux)+        -- 'versionedUniformMap' is only reached via 'isVersionableMap', which+        -- guarantees the body calls no parallel function.+        False+        (worthSequentialising map_lam)+        (segments <> pure w)+        full_body+        outer_body+        intra'++  pure $ insertRegulars (map distResTag ress) match_res env++flattenUniformRedomap ::+  FlattenOps ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  Pat Type ->+  StmAux () ->+  SubExp ->+  [VName] ->+  ScremaForm SOACS ->+  [Reduce SOACS] ->+  Lambda SOACS ->+  FlattenM DistEnv+flattenUniformRedomap ops segments env inps res pat aux w arrs form reds map_lam = do+  let outer_only = transformUniformRedomap (flattenSegLevel ops) segments env inps w arrs reds map_lam+  gpu_scope <- askScope+  let pp_scope = castScope $ scopeOfDistInputs inps <> gpu_scope+  factored <- factorScremaForParallelism funHasParallelism pp_scope (stmAuxCerts aux) pat w arrs form+  case factored of+    Just body ->+      versionScanRed ops "uniform_redomap_alt" segments env inps res aux w body outer_only+    Nothing -> do+      elems' <- outer_only+      pure $ insertRegulars (map distResTag res) elems' env+  where+    funHasParallelism = flattenFunHasParallelism ops++flattenSegRedomap ::+  FlattenOps ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  SubExp ->+  [VName] ->+  [Reduce SOACS] ->+  Lambda SOACS ->+  FlattenM DistEnv+flattenSegRedomap ops segments env inps res w arrs reds map_lam = do+  reps <- mapM (segOpInputRep lvl segments env inps) arrs+  let sing_red = singleReduce reds+      hasNoFreeVariant = allNames (not . isVariant inps . Var) (freeIn sing_red <> freeIn map_lam)+  (ws_F, ws_O, ws_S, elems, elems_kind) <-+    prepareSegOpInputs lvl segments env inps w reps arrs hasNoFreeVariant+  nes' <- mapM (readNeutral segments env inps) (redNeutral sing_red)+  let sing_red' = sing_red {redNeutral = nes'}+  let free = freeIn map_lam+  free_and_sizes <- freeWithTypeDeps inps free+  ws <- dataArr lvl segments env inps w+  (_, _, ws_data) <- doRepIota lvl ws_S+  (free_replicated, replicated) <-+    fmap unzip . sequence $+      mapMaybe+        (onMapFreeVar lvl segments env inps ws (ws_F, ws_O, ws_data))+        free_and_sizes+  (free_env, free_inputs) <- mapArraysToInputs free_replicated replicated++  new_segment <- arraySize 0 <$> lookupType ws_F+  let readFree is = readInputs [new_segment] free_env is free_inputs+  (red_elems, mapout_elems) <-+    genSegRedomap lvl ws_S ws_F ws_O elems sing_red' (soacsLambdaToGPU map_lam) readFree+  red_elems' <- forM red_elems $ \v -> do+    v_t <- lookupType v+    letExp (baseName v <> "_reshaped") . BasicOp $+      Reshape v $+        reshapeAll (arrayShape v_t) (segmentsShape segments)+  let (red_res, map_res) = splitAt (redResults reds) res+  env' <-+    insertSegOpMapResults+      segments+      ws_S+      ws_F+      ws_O+      elems_kind+      (zip map_res mapout_elems)+      env+  pure $ insertRegulars (map distResTag red_res) red_elems' env'+  where+    lvl = flattenSegLevel ops++flattenUniformMaposcanomap ::+  FlattenOps ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  Pat Type ->+  StmAux () ->+  SubExp ->+  [VName] ->+  ScremaForm SOACS ->+  [Scan SOACS] ->+  Lambda SOACS ->+  Lambda SOACS ->+  FlattenM DistEnv+flattenUniformMaposcanomap ops segments env inps res pat aux w arrs form scans post_lam map_lam = do+  let outer_only =+        transformUniformMaposcanomap lvl segments env inps w arrs scans post_lam map_lam+  gpu_scope <- askScope+  let pp_scope = castScope $ scopeOfDistInputs inps <> gpu_scope+  factored <- factorScremaForParallelism funHasParallelism pp_scope (stmAuxCerts aux) pat w arrs form+  case factored of+    Just body ->+      versionScanRed ops "uniform_maposcanomap_alt" segments env inps res aux w body outer_only+    Nothing -> do+      elems' <- outer_only+      pure $ insertRegulars (map distResTag res) elems' env+  where+    funHasParallelism = flattenFunHasParallelism ops+    lvl = flattenSegLevel ops++flattenSegMaposcanomap ::+  FlattenOps ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  SubExp ->+  [VName] ->+  [Scan SOACS] ->+  Lambda SOACS ->+  Lambda SOACS ->+  FlattenM DistEnv+flattenSegMaposcanomap ops segments env inps res w arrs scans post_lam map_lam = do+  reps <- mapM (segOpInputRep lvl segments env inps) arrs+  let hasNoFreeVariant = allNames (not . isVariant inps . Var) (freeIn post_lam <> freeIn map_lam <> foldMap freeIn scans)+  (ws_F, ws_O, ws_S, elems, elems_kind) <-+    prepareSegOpInputs lvl segments env inps w reps arrs hasNoFreeVariant+  let free = freeIn map_lam <> freeIn post_lam+  free_and_sizes <- freeWithTypeDeps inps free+  ws <- dataArr lvl segments env inps w+  (_, _, ws_data) <- doRepIota lvl ws_S+  (free_replicated, replicated) <-+    fmap unzip . sequence $+      mapMaybe+        (onMapFreeVar lvl segments env inps ws (ws_F, ws_O, ws_data))+        free_and_sizes+  (free_env, free_inputs) <- mapArraysToInputs free_replicated replicated+  new_segment <- arraySize 0 <$> lookupType ws_F+  let readFree is = readInputs [new_segment] free_env is free_inputs+  elems' <- doSegMaposcanomap lvl scans ws_F elems post_lam map_lam segments inps env readFree+  insertSegOpMapResults+    segments+    ws_S+    ws_F+    ws_O+    elems_kind+    (zip res elems')+    env+  where+    lvl = flattenSegLevel ops++flattenPlainMap ::+  FlattenOps ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  Pat Type ->+  StmAux () ->+  SubExp ->+  [VName] ->+  Lambda SOACS ->+  FlattenM DistEnv+flattenPlainMap ops segments env inps res pat aux w arrs map_lam = do+  map_res <-+    transformMap ops (stmAuxAttrs aux) segments env inps pat w arrs map_lam+  insertRepsM (zip (map distResTag res) map_res) env++flattenOtherScrema ::+  FlattenOps ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  Pat Type ->+  StmAux () ->+  SubExp ->+  [VName] ->+  ScremaForm SOACS ->+  FlattenM DistEnv+flattenOtherScrema ops segments env inps res pat aux w arrs form = do+  gpu_scope <- askScope+  let pp_scope = castScope $ scopeOfDistInputs inps <> gpu_scope+  factored <- factorScremaForParallelism funHasParallelism pp_scope (stmAuxCerts aux) pat w arrs form+  case factored of+    Just body -> do+      reps <- distributeAndFlattenBody ops segments "factorScremaForParallelism_body" env inps res body+      insertRepsM (zip (map distResTag res) reps) env+    Nothing+      -- XXX: here we silently sequentialise any SOAC that is not handled+      -- above if it is possible to do so. We need to make sure that we+      -- actually handle everything we care about!+      | shouldDissectForm form ->+          error "flattenScrema: complex Screma survived preprocessing"+      | all isRegularDistResult res ->+          flattenScalarStm ops segments env inps res $ Let pat aux (Op (Screma w arrs form))+      | otherwise -> do+          -- XXX: The results are nonuniform, so we cannot run the SOAC+          -- unchanged inside a kernel. Sequentialise it to a loop and+          -- flatten that instead. This does lose us potential parallelism.+          -- A solution would be to preprocess such cases to express them in+          -- terms of loops and maps instead, which we can indeed handle.+          stms <-+            preprocessStms pp_scope+              =<< runSimplifiedBuilder+                pp_scope+                (auxing aux $ FOT.transformSOAC pat $ Screma w arrs form)+          let body = mkBody stms $ varsRes $ patNames pat+          reps <- distributeAndFlattenBody ops segments "sequentialised_soac" env inps res body+          insertRepsM (zip (map distResTag res) reps) env+  where+    funHasParallelism = flattenFunHasParallelism ops++flattenScrema ::+  FlattenOps ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  (Pat Type, StmAux ()) ->+  (SubExp, [VName], ScremaForm SOACS) ->+  FlattenM DistEnv+flattenScrema ops segments env inps res (pat, aux) (w, arrs, form)+  | Just (reds, map_lam) <- isRedomapSOAC form,+    not $ isVariant inps w,+    all isRegularDistResult res,+    all (isRegularInputArr env inps) arrs,+    all (suitableUniformOperator inps . redLambda) reds =+      flattenUniformRedomap ops segments env inps res pat aux w arrs form reds map_lam+  | Just (reds, map_lam) <- isRedomapSOAC form,+    not $ lambdaHasParallelism funHasParallelism map_lam,+    all (\red -> suitableOperator env inps (redLambda red) (redNeutral red)) reds =+      flattenSegRedomap ops segments env inps res w arrs reds map_lam+  | Just (post_lam, scans, map_lam) <- isMaposcanomapSOAC form,+    not $ isVariant inps w,+    all isRegularDistResult res,+    all (isRegularInputArr env inps) arrs,+    all (suitableUniformOperator inps . scanLambda) scans =+      flattenUniformMaposcanomap ops segments env inps res pat aux w arrs form scans post_lam map_lam+  | Just (post_lam, scans, map_lam) <- isMaposcanomapSOAC form,+    not $ lambdaHasParallelism funHasParallelism map_lam,+    not $ lambdaHasParallelism funHasParallelism post_lam,+    all (\scan -> suitableOperator env inps (scanLambda scan) (scanNeutral scan)) scans =+      flattenSegMaposcanomap ops segments env inps res w arrs scans post_lam map_lam+  | Just map_lam <- isMapSOAC form,+    isVersionableMap funHasParallelism lvl inps env w res map_lam =+      versionedUniformMap ops segments env inps res pat aux w arrs map_lam+  | Just map_lam <- isMapSOAC form =+      flattenPlainMap ops segments env inps res pat aux w arrs map_lam+  | otherwise =+      flattenOtherScrema ops segments env inps res pat aux w arrs form+  where+    funHasParallelism = flattenFunHasParallelism ops+    lvl = flattenSegLevel ops++-- Note [FlatMap element counting]+--+-- 'FlatMap' counts in units of the lambda's element type @b@: the data array+-- for irregular results is @[m]b@, the shape and offset arrays hold per-segment+-- counts of @b@-elements, and the flag array has one entry per @b@-element.+--+-- The flattening pass, however, represents irregular data with an+-- 'IrregularRep' whose 'irregularD' is primitive (a flat array of scalars), so+-- its segment sizes and flags count *scalars*. When @b@ is itself a @c@-element+-- array there are @c@ scalars per @b@-element ('flatMapElemsPer'), so the two+-- notions of "size" differ by a factor of @c@.+--+-- The flattening rules therefore convert between the two: the metadata that+-- 'FlatMap' surfaces directly (the total @m@, and the shape, flag, and offset+-- arrays) is put in @b@-units by dividing scalar-unit sizes by @c@. The data+-- arrays themselves, however, keep their natural primitive 'IrregularRep' with+-- scalar-unit structure - the "'irregularD' is always primitive" invariant is+-- preserved. A consumer of such a data array reconstructs the @b@-element rows+-- from the primitive rep: 'onMapIrregularInputArr' does so for a plain map, and+-- 'prepareSegOpInputs' for a segmented redomap/scan.++-- Note [Ill-formed inner-map lambda]+--+-- 'flattenFlatMapNested' hands the inner-map machinery ('transformMap') an+-- ordinary 'Lambda' obtained from the 'FlatMap' 'ExtLambda' by instantiating+-- the existential size with the lambda's own leading result. That lambda is+-- ill-formed: the size is bound by a statement in the body, and so is+-- meaningless in the return type, which describes the lambda from outside.+--+-- We do it because 'transformMap' requires a 'Lambda', whose types cannot+-- express a result whose outer size varies per iteration, which is exactly what+-- an irregular 'FlatMap' result is.+--+-- Nothing downstream derives meaning from the return type - the results are+-- described by 'map_pat' throughout - but it is still traversed in two places,+-- which is what makes this contained rather than harmless:+--+-- - 'distributeAndTransformInnerMap' computes the free variables that must be+--   replicated into the body, and must skip the return type; otherwise the size+--   is looked up in the enclosing scope, where it does not exist.+--+-- - 'preprocessLambda' simplifies the return type outside the scope of the+--   parameters, where a size that resolves to nothing is left alone.+--+-- FIXME: Teaching the inner-map machinery to accept an 'ExtLambda' would remove+-- the need for all of this.++-- | The number of scalars per element of a 'FlatMap' lambda's (first) value+-- result; see Note [FlatMap element counting].+flatMapElemsPer :: ExtLambda SOACS -> TPrimExp Int64 VName+flatMapElemsPer lam =+  product $ map pe64 $ arrayDims $ head $ flatMapRowTypes lam++-- | Divide each entry of a per-segment size (or offset) array by the given+-- factor, converting scalar-unit counts to element-unit counts. A no-op when the+-- factor is statically 1 (a scalar element type). See Note [FlatMap element+-- counting].+scaleSizesDown :: SegLevel -> TPrimExp Int64 VName -> VName -> FlattenM VName+scaleSizesDown _ elems_per arr | elems_per == 1 = pure arr+scaleSizesDown lvl elems_per arr = do+  arr_t <- lookupType arr+  letExp "scaled_sizes" <=< segMap lvl (arrayDims arr_t) $ \gtids -> do+    x <- letSubExp "x" =<< eIndex arr (map eSubExp gtids)+    x_b <- letSubExp "x_b" =<< toExp (pe64 x `div` elems_per)+    pure [subExpRes x_b]++-- | Subsample a scalar-unit flag array to element units by taking every+-- @elems_per@-th entry, giving one flag per element. A no-op when the factor is+-- statically 1. This reuses the (already-built) flag array rather than+-- recomputing segment starts with a segmented scan. See Note [FlatMap element+-- counting].+subsampleFlags :: SegLevel -> TPrimExp Int64 VName -> VName -> FlattenM VName+subsampleFlags _ elems_per arr | elems_per == 1 = pure arr+subsampleFlags lvl elems_per arr = do+  big_m <- arraySize 0 <$> lookupType arr+  m <- letSubExp "flags_m" =<< toExp (pe64 big_m `div` elems_per)+  letExp "elem_flags" <=< segMap lvl (MkSolo m) $ \(MkSolo i) -> do+    flag <- letSubExp "flag" =<< eIndex arr [toExp $ pe64 i * elems_per]+    pure [subExpRes flag]++-- | Flattening rule for a top-level 'FlatMap', which is a very thin wrapper+-- over just flattening the lambda. The result produced by 'FlatMap' corresponds+-- exactly to the internal irregular representation, so we simply distribute the+-- lambda over the @w@ segments and obtain an 'IrregularRep' for each result+-- (they share the same segment structure, as required by the type of+-- 'FlatMap').+flattenFlatMap ::+  FlattenOps ->+  Pat Type ->+  SubExp ->+  [VName] ->+  ExtLambda SOACS ->+  FlattenM ()+flattenFlatMap ops pat w arrs lam = do+  let segments = [w]+      inps =+        zipWith+          (\p arr -> (paramName p, DistInputFree arr (paramType p)))+          (lambdaParams lam)+          arrs+      elem_ts = flatMapRowTypes lam+      (m_name, s_name, f_name, o_name, value_names) =+        case patNames pat of+          (a : b : c : d : hs) -> (a, b, c, d, hs)+          _ -> error "flattenFlatMap: pattern too short"+      (d_names, r_names) = flatMapSplitValues lam value_names+      -- A nonuniform result is distributed as a variably sized array, a uniform+      -- one as an ordinary array with one element per segment.+      distTypeOf t+        | flatMapNonuniform t = DistType segments (Rank 1) $ static $ rowType t+        | otherwise = DistType segments (Rank 0) $ static t+      res =+        zipWith3+          (\i v t -> DistResult (ResTag i) (distTypeOf t) v)+          [0 ..]+          value_names+          value_ts+      -- The lambda's leading size result is of no use here, as the metadata is+      -- derived from the segment structure of the flattened irregular results.+      body = lambdaBody lam+      body' = body {bodyResult = drop 1 (bodyResult body)}+  reps <- distributeAndFlattenBody ops segments "flatmap" mempty inps res body'+  let (irreg_reps, reg_reps) = flatMapSplitValues lam reps+  irregs <- forM irreg_reps $ \case+    Irregular ir -> ensureDenseIrregular lvl "flatmap_res" ir+    Regular _ -> error "flattenFlatMap: irregular result is not irregular"+  -- The regular results (outer size @w@) are already in the form we want.+  forM_ (zip r_names reg_reps) $ \(v, rep) -> case rep of+    Regular v' -> letBindNames [v] $ BasicOp $ SubExp $ Var v'+    Irregular _ -> error "flattenFlatMap: regular result is not regular"+  case (irregs, elem_ts) of+    (ir0 : _, _ : _) -> do+      -- The flattening structure arrays are in units of scalars, but the source+      -- 'flatmap' wants them in units of the lambda's element type @b@, of+      -- which there are 'elems_per' scalars. All results share the same segment+      -- structure, so we derive the source metadata once, from the first+      -- result:+      --+      --   * The shape and offset arrays (per-segment, size @w@) are the+      --     scalar-unit 'irregularS' and 'irregularO' scaled down by+      --     'elems_per'.+      --+      --   * the flag array (per-element, size @m@) is 'irregularF' subsampled+      --     with stride 'elems_per'.+      let elems_per = flatMapElemsPer lam+      -- The total size is the number of data elements.+      big_m <- arraySize 0 <$> lookupType (irregularD ir0)+      m <- letSubExp "flatmap_m" =<< toExp (pe64 big_m `div` elems_per)+      letBindNames [m_name] $ BasicOp $ SubExp m+      s <- scaleSizesDown lvl elems_per (irregularS ir0)+      o <- scaleSizesDown lvl elems_per (irregularO ir0)+      flags <- subsampleFlags lvl elems_per (irregularF ir0)+      -- The per-segment metadata (outer size @w@): shape and offset arrays.+      bindReshape s_name s (Shape [w])+      bindReshape o_name o (Shape [w])+      -- The per-element flag array (outer size @m@).+      bindReshape f_name flags (Shape [Var m_name])+      -- The concatenated data arrays, reshaped from the flat scalar data to+      -- @[m]b@ (the outer dimension being the freshly bound total size).+      forM_ (zip3 d_names irregs elem_ts) $ \(v, ir, elem_t) ->+        bindReshape v (irregularD ir) (Shape [Var m_name] <> Shape (arrayDims elem_t))+    _ -> error "flattenFlatMap: FlatMap with no irregular results"+  where+    lvl = flattenSegLevel ops+    value_ts = drop 1 (lambdaReturnType lam)+    static =+      fromMaybe (error "flattenFlatMap: existential size.") . hasStaticShape+    bindReshape name arr newshape = do+      arr_t <- lookupType arr+      letBindNames [name] . BasicOp $+        Reshape arr (reshapeAll (arrayShape arr_t) newshape)++-- | Flattening rule for a 'FlatMap' nested inside an enclosing map-nest. A+-- 'FlatMap' is just a nonuniform map with implicit concatenation, so we run it+-- through the ordinary inner-map machinery, which already produces - for a+-- variably-sized result - an 'IrregularRep' whose segment sizes are the+-- per-enclosing-segment concatenated lengths and whose data is the+-- concatenation. That directly gives the data arrays (results 4..) and, as+-- their per-enclosing segment sizes, the total sizes @m@ (result 0).+--+-- The metadata is derived from the shape array (result 1) - the per-iteration+-- output sizes - which the inner-map does not itself surface. The lambda+-- returns exactly that size as its leading result, so mapping it over the nest+-- yields the irregular array of per-iteration sizes, segmented by the 'FlatMap'+-- width. From it we then compute the offset array (result 3), as the+-- per-enclosing exclusive prefix sum, and the flag array (result 2).+--+-- See also Note [FlatMap element counting].+flattenFlatMapNested ::+  FlattenOps ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  StmAux () ->+  SubExp ->+  [VName] ->+  ExtLambda SOACS ->+  FlattenM DistEnv+flattenFlatMapNested ops segments env inps res aux w arrs lam = do+  size_name <- newVName "flatmap_sizes"+  data_names <- mapM (const $ newVName "flatmap_data") val_ts+  let size_pe = PatElem size_name $ arrayOfRow (Prim int64) w+      data_pes = zipWith (\v t -> PatElem v $ arrayOfRow t w) data_names val_ts+      map_pat = Pat $ size_pe : data_pes+  let lvl = flattenSegLevel ops+  certifying (distCerts inps aux env) $ do+    reps <- transformMap ops (stmAuxAttrs aux) segments env inps map_pat w arrs map_lam+    case reps of+      shape_rep : data_reps+        | Just first_data <- firstIrregular data_reps -> do+            -- An 'IrregularRep' always stores primitive data, so for a+            -- non-scalar element type @b@ the data result's segment sizes are+            -- in units of scalars, of which there are 'elems_per' per+            -- b-element. The source 'flatmap' counts b-elements, so we scale+            -- the affected metadata down.+            let elems_per = flatMapElemsPer lam++            -- The number of b-elements per enclosing segment: the per-segment+            -- totals (result 0), and the segment sizes of the flag array. These+            -- are the data result's segment sizes, scaled to b-element units.+            -- When the data is regular (the fully uniform case), it is the+            -- constant inner size, replicated across the segments.+            s_flag <- case first_data of+              Irregular ir -> scaleSizesDown lvl elems_per (irregularS ir)+              Regular d -> do+                d_t <- lookupType d+                let inner = product $ map pe64 $ drop (segmentsRank segments) $ arrayDims d_t+                n_elem <- letSubExp "flatmap_n_elem" =<< toExp (inner `div` elems_per)+                letExp "flatmap_n" (BasicOp $ Replicate (segmentsShape segments) n_elem)+            let n_rep = Regular s_flag++            -- The offset array (result 3) has the same structure as the shape+            -- array; its values are the exclusive prefix sum of the shape+            -- within each enclosing segment. We also need the shape as a single+            -- flat array of all per-iteration sizes (for the flags).+            (offset_rep, shape_flat) <- case shape_rep of+              Regular sd -> do+                off <- genExPrefixSum lvl "flatmap_offset" sd+                sd_t <- lookupType sd+                n_shape <- letSubExp "flatmap_shape_n" =<< toExp (product $ map pe64 $ arrayDims sd_t)+                flat <-+                  letExp "flatmap_shape_flat" . BasicOp $+                    Reshape sd (reshapeAll (arrayShape sd_t) (Shape [n_shape]))+                pure (Regular off, flat)+              Irregular s_ir -> do+                inc <- genSegPrefixSum lvl "flatmap_offset_inc" (irregularF s_ir) (irregularD s_ir)+                n_off <- arraySize 0 <$> lookupType inc+                off_D <- letExp "flatmap_offset_D" <=< segMap lvl (MkSolo n_off) $ \(MkSolo i) -> do+                  a <- letSubExp "a" =<< eIndex inc [eSubExp i]+                  b <- letSubExp "b" =<< eIndex (irregularD s_ir) [eSubExp i]+                  off <- letSubExp "off" =<< toExp (pe64 a - pe64 b)+                  pure [subExpRes off]+                pure (Irregular s_ir {irregularD = off_D}, irregularD s_ir)++            -- The flag array (result 2): the segment-start flags over all+            -- b-elements. Every enclosing-segment boundary is also a segment start,+            -- so a single 'doRepIota' over all per-iteration sizes yields the flags+            -- for all enclosing segments at once. Its per-enclosing segment sizes+            -- are @s_flag@ (in b-element units); we build the enclosing flag/offset+            -- structure from those, respecting the invariant that in an+            -- 'IrregularRep' the flag and data arrays share their shape (rather than+            -- borrowing the scalar-unit structure of the data rep).+            (flag_F, flag_O, _) <- doRepIota lvl s_flag+            (flag_D0, _, _) <- doRepIota lvl shape_flat+            -- 'flag_D0' and 'flag_F' both have the total number of b-elements as+            -- their size, but computed by different means; coerce so the rep's flag+            -- and data arrays share a size (the 'IrregularRep' invariant).+            m_flag <- arraySize 0 <$> lookupType flag_F+            flag_D0_t <- lookupType flag_D0+            flag_D <-+              letExp "flatmap_flag_D" . BasicOp $+                Reshape flag_D0 (reshapeAll (arrayShape flag_D0_t) (Shape [m_flag]))+            let flag_rep = Irregular $ IrregularRep s_flag flag_F flag_O flag_D Dense++            -- The data arrays (results 4..). The inner-map machinery already+            -- produced, for each, an 'IrregularRep' with primitive+            -- (scalar-unit) data and matching scalar-unit structure arrays.+            -- These are exactly the reps for the source @[m]b@ data arrays, so+            -- we pass them through unchanged; any consumer reconstructs the+            -- @b@-element rows from the primitive rep via the ordinary+            -- irregular-input machinery. See Note [FlatMap element counting].+            let all_reps = n_rep : shape_rep : flag_rep : offset_rep : data_reps+            insertRepsM (zip (map distResTag res) all_reps) env+      _ -> error "flattenFlatMapNested: FlatMap with no irregular results"+  where+    -- The metadata is derived from a nonuniform result; they all have the same+    -- segment structure, so any of them will do.+    firstIrregular =+      fmap snd+        . find (flatMapNonuniform . fst)+        . zip (drop 1 (lambdaReturnType lam))+    -- This lambda is deliberately ill-formed; see Note [Ill-formed inner-map+    -- lambda].+    map_lam =+      Lambda+        { lambdaParams = lambdaParams lam,+          lambdaReturnType = Prim int64 : val_ts,+          lambdaBody = lambdaBody lam+        }+    -- Only works because we allow just a single Ext in the return type.+    val_ts =+      runIdentity . instantiateShapes (const (pure k)) . drop 1 $+        lambdaReturnType lam+      where+        -- The size of the lambda's irregular results, which it returns before+        -- them, and which is what 'Ext 0' stands for.+        k = case bodyResult (lambdaBody lam) of+          r : _ -> resSubExp r+          [] -> error "flattenFlatMapNested: FlatMap with no results"++-- | Remove certificates that refer to variables free in the lambda (recursing+-- into nested bodies). Such certificates arise on pure operator lambdas (e.g.+-- the combining function of a 'hist' or reduce) through conservative+-- certificate propagation. They are redundant, and cannot be preserved when the+-- operator is lifted into a segmented operation, as the certificate's binding+-- does not survive into the generated kernel.+stripFreeCerts :: Lambda SOACS -> Lambda SOACS+stripFreeCerts lam = lam {lambdaBody = onBody (lambdaBody lam)}+  where+    frees = freeIn lam+    onBody body = body {bodyStms = onStm <$> bodyStms body}+    onStm (Let pat dec e) = Let pat (onDec dec) (onExp e)+    onDec dec =+      dec {stmAuxCerts = Certs $ filter (`notNameIn` frees) $ unCerts $ stmAuxCerts dec}+    onExp = runIdentity . mapExpM mapper+    mapper = identityMapper {mapOnBody = const (pure . onBody)}++flattenHist ::+  FlattenOps ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  (Pat Type, StmAux ()) ->+  (SubExp, [VName], [Futhark.IR.SOACS.HistOp SOACS], Lambda SOACS) ->+  FlattenM DistEnv+flattenHist ops segments env inps res (_pat, aux) (w, hist_inputs, hist_ops0, bucket_fun) = do+  -- The operator is a pure combining function; any certificates on its+  -- statements (from conservative propagation) are redundant and cannot be+  -- preserved when it is lifted into a segmented operation, so drop the ones+  -- referring to variables free in the operator.+  let hist_ops = do+        op <- hist_ops0+        pure $+          op+            { Futhark.IR.SOACS.histOp = stripFreeCerts $ Futhark.IR.SOACS.histOp op+            }+  -- TODO: check for suitableUniformOperator.+  let nonuniform =+        not (all (isRegularInputArr env inps) hist_inputs)+          || isVariant inps w+          || not (all isRegularDistResult res)+  if nonuniform+    then do+      gpu_scope <- askScope+      let scope = castScope $ scopeOfDistInputs inps <> gpu_scope+      (hist_res, stms) <-+        runBuilderT+          ( auxing aux $+              doHist "nonuniform_hist" hist_ops hist_inputs $ \params ->+                map resSubExp <$> eLambda bucket_fun (map eParam params)+          )+          scope+      let body = mkBody stms $ varsRes $ concat hist_res+      reps <- distributeAndFlattenBody ops segments "non_uniform_hist_body" env inps res body+      insertRepsM (zip (map distResTag res) reps) env+    else do+      let new_segment = segments <> pure w+      lifted_inps <- forM hist_inputs $ \hist_inp -> do+        t <- lookupInputType inps hist_inp+        let expectedShape = segmentsShape segments <> arrayShape t+        liftSubExpRegular lvl segments inps env expectedShape (Var hist_inp)+      hist_ops' <- forM hist_ops $ \(Futhark.IR.SOACS.HistOp num_bins rf dests nes op) -> do+        nes' <- mapM (readNeutral segments env inps) nes+        let rr (DistType _ _ t) = t+        let ts = map (rr . distResType) res+        let expectedShapes = map (\t -> segmentsShape segments <> arrayShape t) ts+        dests' <- mapM (\(shape, var) -> liftSubExpRegular lvl segments inps env shape (Var var)) (zip expectedShapes dests)+        pure $ Futhark.IR.SOACS.HistOp num_bins rf dests' nes' op+      let free = freeIn bucket_fun+      let isDest = flip elem $ concatMap Futhark.IR.SOACS.histDest hist_ops'+          free_notDest = filter (not . isDest) (namesToList free)+      free_and_sizes <- freeWithTypeDeps inps (namesFromList free_notDest)+      (free_replicated, replicated) <-+        fmap unzip . sequence $+          mapMaybe+            (onMapFreeVarMultiDim lvl segments w env inps)+            free_and_sizes+      (free_env, free_inputs) <- mapArraysToInputs free_replicated replicated+      let readFree is = readInputs new_segment free_env is free_inputs+      hist_res <-+        certifying (distCerts inps aux env) $+          genUniformSegHist lvl "Uniform_segHist" new_segment hist_ops' (soacsLambdaToGPU bucket_fun) lifted_inps readFree+      pure $ insertRegulars (map distResTag res) hist_res env+  where+    lvl = flattenSegLevel ops
+ src/Futhark/Pass/Flatten/WithAcc.hs view
@@ -0,0 +1,322 @@+-- | Flattening of 'WithAcc'.+--+-- The basic idea is that in the nonuniform case, we change the 'WithAcc' to be+-- over the data array of the irregular representation. We then update all the+-- 'UpdateAcc' operations to compute flat indexes, via the usual metadata+-- arrays.+module Futhark.Pass.Flatten.WithAcc+  ( flattenWithAcc,+  )+where++import Control.Monad+import Data.Foldable+import Data.List qualified as L+import Data.Map qualified as M+import Futhark.IR.GPU+import Futhark.IR.SOACS+import Futhark.MonadFreshNames+import Futhark.Pass.Flatten.Distribute+import Futhark.Pass.Flatten.General+import Futhark.Tools+import Futhark.Transform.ToGPU (soacsLambdaToGPU)+import Prelude hiding (div, rem)++indexIrreg ::+  (MonadBuilder m, BranchType (Rep m) ~ ExtType) =>+  Segments ->+  DistEnv ->+  IrregularRep ->+  SubExp ->+  Safety ->+  ShapeBase SubExp ->+  [SubExp] ->+  m SubExp+indexIrreg _segments _env rep is safety shape js = do+  offset <- letSubExp "uacc_segment_offset" =<< eIndex (irregularO rep) [eSubExp is]+  flat <-+    letSubExp "flat_uacc_idx" <=< toExp $+      pe64 offset+        + flattenIndex+          (map pe64 (shapeDims shape))+          (map pe64 js)+  case safety of+    Unsafe -> pure flat+    -- The original per-segment index 'js' may be out of bounds, which for a+    -- safe 'UpdateAcc' is a deliberate no-op. Adding the segment offset might+    -- turn such an out-of-bounds index into an in-bounds index into a+    -- neighbouring segment, so we must preserve the out-of-boundedness+    -- explicitly.+    Safe -> do+      inbounds <- letSubExp "uacc_inbounds" =<< eShapeInBounds shape (map eSubExp js)+      letSubExp "uacc_flat_idx"+        =<< eIf+          (eSubExp inbounds)+          (eBody [eSubExp flat])+          (eBody [eSubExp (intConst Int64 (-1))])++-- If just one input is nonuniform, we treat them all as nonuniform.+flattenWithAcc ::+  FlattenOps ->+  Segments ->+  DistEnv ->+  DistInputs ->+  [DistResult] ->+  Pat Type ->+  StmAux () ->+  [WithAccInput SOACS] ->+  Lambda SOACS ->+  FlattenM DistEnv+flattenWithAcc ops segments env inps distres _withacc_pat withacc_aux withacc_inputs acc_lam = do+  lam_params' <- newAccLamParams $ lambdaParams acc_lam++  iota_w <- genShapeIota (flattenSegLevel ops) $ segmentsShape segments++  iota_p <- newParam "iota_p" $ Prim int64+  -- Type in DistInputFree is parameter type+  let iota_w_t = Prim int64+  let iota_se = Var (paramName iota_p)++  -- Potentially change to distres option.+  let nonuniform =+        not $+          all (\(_, arrs, _) -> all (isRegularInputArr env inps) arrs) withacc_inputs++  (withacc_inputs', trAccIndex, non_uniform_reps) <-+    if nonuniform+      then do+        (withacc_inputs', input_reps) <-+          mapAndUnzipM onNonuniformInput withacc_inputs+        let trAccIndex c safety is = do+              ((shape, _, _), rep : _) <-+                L.lookup c $+                  zip (map paramName lam_params') $+                    zip withacc_inputs input_reps+              Just $ L.singleton <$> indexIrreg segments env rep iota_se safety shape is+        pure (withacc_inputs', trAccIndex, concat input_reps)+      else do+        withacc_inputs' <- mapM onUniformInput withacc_inputs+        let trAccIndex c _safety is = do+              _ <- L.lookup c $ zip (map paramName lam_params') withacc_inputs+              Just $ do+                iota_se_unflat <-+                  mapM (letSubExp "iota_idx" <=< toExp) $+                    unflattenIndex (segmentDims segments) (pe64 iota_se)+                pure $ iota_se_unflat ++ is+        pure (withacc_inputs', trAccIndex, [])+  let trAccShape c = do+        (ispace, _, _) <- L.lookup c $ zip (map paramName lam_params') withacc_inputs'+        pure ispace+      sf = (trAccShape, trAccIndex)++  acc_lam_body <-+    runBodyBuilder $+      localScope (scopeOfLParams lam_params') $+        bodyBind . lambdaBody =<< trLam sf acc_lam++  scope <- askScope+  let orig_acc_params = drop num_accs $ lambdaParams acc_lam+      lam_params_tr = map (trParam sf) lam_params'+      acc_params_tr = drop num_accs lam_params_tr+      interchanged_inps =+        (paramName iota_p, DistInputFree iota_w iota_w_t)+          : [ (paramName p, DistInputFree (paramName acc) (paramType acc))+            | -- This could potentially be wrong but since it's acc type it should be fine.+              (p, acc) <- zip orig_acc_params acc_params_tr+            ]+          ++ inps++  let (withacc_new_inputs, withacc_dstms) =+        distributeBodyWith ops scope segments interchanged_inps acc_lam_body++  withacc_lam' <- mkLambda (map (trParam sf) lam_params') $ do+    env' <- foldM (flattenDistStm ops segments) env withacc_dstms+    reps <-+      mapM+        (liftWithAccResult (flattenSegLevel ops) segments withacc_new_inputs env')+        (zip distres (bodyResult $ lambdaBody acc_lam))+    concat <$> mapM repToResults reps++  withacc_out_vs <-+    certifying (distCerts inps withacc_aux env) $+      letTupExp "withacc_flatten_out" (WithAcc withacc_inputs' withacc_lam')++  -- The accumulator results are handled differently in the nonuniform case,+  -- since we do not have metadata for them and since all of them are turned+  -- flat even when they might actually be regular. We can still turn the+  -- distres that are regular into Regulars here.+  let num_acc_results = sum [length arrs | (_, arrs, _) <- withacc_inputs]+      (withacc_out_vs_wo, withacc_out_vs_no) = splitAt num_acc_results withacc_out_vs+      (distres_withacc, distres_normal) = splitAt num_acc_results distres++  let out_reps_normal = resultToResRepsByDistResult distres_normal withacc_out_vs_no+  out_reps_withacc <-+    if nonuniform+      then mapM mkNonuniformWithAccRep (zip3 withacc_out_vs_wo non_uniform_reps distres_withacc)+      else pure $ map Regular withacc_out_vs_wo+  insertRepsM (zip (map distResTag $ distres_withacc ++ distres_normal) (out_reps_withacc ++ out_reps_normal)) env+  where+    newAccLamParams ps = do+      let (cert_ps, acc_ps) = splitAt num_accs ps+      -- Should not rename the certificates.+      acc_ps' <- forM acc_ps $ \(Param attrs v t) ->+        Param attrs <$> newName v <*> pure t+      pure $ cert_ps <> acc_ps'++    num_accs = length withacc_inputs++    onOpWithIndexRank index_rank (op_lam, nes) = do+      -- We need to add an additional index parameter because we are extending+      -- the index space of the accumulator. In the uniform case we have the+      -- full index space of the segments, while in the nonuniform case we only+      -- have one additional dimension.+      idx_ps <- replicateM index_rank $ newParam "idx" $ Prim int64+      pure+        ( soacsLambdaToGPU $+            op_lam {lambdaParams = idx_ps <> lambdaParams op_lam},+          nes+        )++    -- Let's use liftSubExpRegular here+    onUniformInput (shape, arrs, op) =+      (segmentsShape segments <> shape,,)+        <$> mapM onArr arrs+        <*> traverse (onOpWithIndexRank (segmentsRank segments)) op+      where+        onArr arr = do+          arr_t <- lookupInputType inps arr+          let arr_shape = arrayShape arr_t+              expected_shape = segmentsShape segments <> arr_shape+          liftSubExpRegular (flattenSegLevel ops) segments inps env expected_shape (Var arr)++    onNonuniformOp rank (op, nes) = do+      let (old_idx_ps, value_ps) = splitAt rank $ lambdaParams op+          old_indices_used =+            any (\p -> paramName p `nameIn` freeIn (lambdaBody op)) old_idx_ps+      -- XXX: It is possible to change the lambda body to restore the indices based on flat_idx_p and handle this.+      when old_indices_used $+        error "flattenWithAcc: accumulator operator uses nonuniform indices"+      flat_idx_p <- newParam "flat_idx" $ Prim int64+      pure+        ( soacsLambdaToGPU $+            op+              { lambdaParams = flat_idx_p : value_ps+              },+          nes+        )+    onNonuniformInput (shape, arrs, op) = do+      reps <- mapM (getIrregRep (flattenSegLevel ops) segments env inps) arrs+      -- We need to ensure that the irregular arrays are dense+      reps_dense <- mapM (ensureDenseIrregular (flattenSegLevel ops) "withacc_input") reps+      let arrs' = map irregularD reps_dense+      w <- fmap (arraySize 0) . lookupType $ head arrs'+      -- We need to reshape to make sure all of the inputs have the same shape+      arrs'' <- forM arrs' $ \v -> do+        v_t <- lookupType v+        letExp (baseName v <> "_withacc_input_reshaped") . BasicOp $+          Reshape v $+            reshapeAll (arrayShape v_t) (Shape [w])+      (,reps_dense) . (Shape [w],arrs'',) <$> traverse (onNonuniformOp (shapeRank shape)) op++    -- The irregular kind is not carried through the results of the+    -- WithAcc, and 'mkNormalResReps' reconstructs the rep as Dense, so+    -- any irregular rep must actually be made dense before it crosses+    -- the WithAcc boundary.+    ensureDenseRep lvl (Irregular irreg) =+      Irregular <$> ensureDenseIrregular lvl "withacc_result" irreg+    ensureDenseRep _ rep = pure rep++    liftWithAccResult lvl segs inputs env' (dist_res, res) =+      case resSubExp res of+        Var v -> do+          let (Just (t, rep)) = M.lookup v $ inputReps inputs env'+          if isAcc t+            then+              pure rep+            else+              ensureDenseRep lvl =<< liftDistResultRep lvl segs inputs env' dist_res res+        Constant _ -> ensureDenseRep lvl =<< liftDistResultRep lvl segs inputs env' dist_res res++    repToResults (Regular v) =+      pure [SubExpRes mempty $ Var v]+    repToResults (Irregular irreg) =+      map (SubExpRes mempty . Var) <$> irregResults irreg++    irregResults+      ( IrregularRep+          { irregularS = segs,+            irregularF = flags,+            irregularO = offsets,+            irregularD = elems+          }+        ) = do+        flags_t <- lookupType flags+        t <- lookupType elems+        num_data <- letExp "num_data" =<< toExp (product $ map pe64 $ arrayDims t)+        let shape = Shape [Var num_data]+        flags' <- letExp "flags" $ BasicOp $ Reshape flags $ reshapeAll (arrayShape flags_t) shape+        elems' <- letExp "elems" $ BasicOp $ Reshape elems $ reshapeAll (arrayShape t) shape+        pure [num_data, segs, flags', offsets, elems']++    mkNonuniformWithAccRep (v, rep, dist_res)+      | isRegularDistResult dist_res = do+          let DistType _ _ t = distResType dist_res+              expectedShape = segmentsShape segments <> arrayShape t+          v_t <- lookupType v+          v_reshaped <-+            letExp "actual_regular_with_acc_res" . BasicOp $+              Reshape v (reshapeAll (arrayShape v_t) expectedShape)+          pure $ Regular v_reshaped+      | otherwise =+          pure $ Irregular $ rep {irregularD = v}++    trType ::+      (VName -> Maybe Shape, VName -> Safety -> [SubExp] -> Maybe (Builder SOACS [SubExp])) ->+      TypeBase shape u ->+      TypeBase shape u+    trType sf (Acc acc _ ts u)+      | Just shape <- fst sf acc =+          Acc acc shape ts u+    trType _ t = t++    trParam ::+      (VName -> Maybe Shape, VName -> Safety -> [SubExp] -> Maybe (Builder SOACS [SubExp])) ->+      Param (TypeBase Shape u) ->+      Param (TypeBase Shape u)+    trParam sf = fmap $ trType sf++    trBody sf (Body dec stms res) =+      Body dec <$> collectStms_ (traverse_ onStm stms) <*> pure res+      where+        onStm (Let pat aux e) =+          addStm . Let (fmap (trType sf) pat) aux =<< trExp sf pat e++    trLam sf (Lambda params ret body) =+      Lambda (map (trParam sf) params) (map (trType sf) ret) <$> trBody sf body++    trSOAC sf = mapSOACM mapper+      where+        mapper =+          identitySOACMapper {mapOnSOACLambda = trLam sf}++    trExp sf _ (WithAcc acc_inputs lam) =+      WithAcc acc_inputs <$> trLam sf lam+    trExp sf (Pat [PatElem _ acc_t]) (BasicOp (UpdateAcc safety acc is ses)) = do+      case acc_t of+        Acc cert _ _ _+          | Just mk <- snd sf cert safety is -> do+              is' <- mk+              pure $ BasicOp $ UpdateAcc safety acc is' ses+        _ ->+          pure $ BasicOp $ UpdateAcc safety acc is ses+    trExp sf _ e = mapExpM mapper e+      where+        mapper =+          identityMapper+            { mapOnBody = \_ -> trBody sf,+              mapOnRetType = pure . trType sf,+              mapOnBranchType = pure . trType sf,+              mapOnFParam = pure . trParam sf,+              mapOnLParam = pure . trParam sf,+              mapOnOp = trSOAC sf+            }
+ src/Futhark/Pass/NoGrid.hs view
@@ -0,0 +1,35 @@+-- | Remove grid information from all SegOp operations - the idea is to remove+-- all GPU-specific information from a GPU IR program.+module Futhark.Pass.NoGrid (noGrid) where++import Control.Monad.Identity+import Futhark.IR.GPU+import Futhark.Pass++noGrid :: Pass GPU GPU+noGrid =+  Pass "no-grid" "remove grid info from SegOps" $+    intraproceduralTransformation optimise+  where+    optimise _scope stms = pure $ fmap onStm stms+    onStm (Let pat aux e) =+      Let pat aux $+        mapExp+          (identityMapper @GPU)+            { mapOnOp = pure . onOp,+              mapOnBody = const $ pure . onBody,+              mapOnFParam = pure+            }+          e+    onBody body = body {bodyStms = fmap onStm (bodyStms body)}++    onLevel (SegThread v _) = SegThread v Nothing+    onLevel (SegBlock v _) = SegBlock v Nothing+    onLevel (SegThreadInBlock v) = SegThreadInBlock v++    onOp (SegOp op) =+      SegOp . runIdentity $+        mapSegOpM+          (identitySegOpMapper {mapOnSegOpLevel = pure . onLevel})+          op+    onOp op = op
src/Futhark/Passes.hs view
@@ -38,9 +38,9 @@ import Futhark.Pass.ExplicitAllocations.GPU qualified as GPU import Futhark.Pass.ExplicitAllocations.MC qualified as MC import Futhark.Pass.ExplicitAllocations.Seq qualified as Seq-import Futhark.Pass.ExtractKernels import Futhark.Pass.ExtractMulticore import Futhark.Pass.FirstOrderTransform+import Futhark.Pass.Flatten import Futhark.Pass.LiftAllocations as LiftAllocations import Futhark.Pass.LowerAllocations as LowerAllocations import Futhark.Pass.Simplify@@ -84,9 +84,11 @@ gpuPipeline :: Pipeline SOACS GPU gpuPipeline =   standardPipeline-    >>> onePass extractKernels+    >>> onePass flattenSOACs     >>> passes       [ simplifyGPU,+        -- For getting rid of builtins added by flattening.+        removeDeadFunctions,         addGlobalParams,         optimiseGenRed,         simplifyGPU,
src/Futhark/Tools.hs view
@@ -7,12 +7,14 @@     redomapToMapAndReduce,     scanomapToMapAndScan,     maposcanomapToMapScanAndMap,+    maposcanomapToMaposcanAndMap,     dissectScrema,     extractPostLambda,     sequentialStreamWholeArray,     partitionChunkedFoldParameters,     withAcc,     doScatter,+    doHist,     addBinOp,     addLambda, @@ -28,7 +30,7 @@ import Futhark.Construct import Futhark.IR import Futhark.IR.SOACS.SOAC-import Futhark.Util (mapAccumLM)+import Futhark.Util (chunks, mapAccumLM)  splitScanOrRedomap ::   (MonadFreshNames m) =>@@ -152,6 +154,30 @@   where     tempRes res = newIdent "temp_res" $ res `arrayOfRow` w +maposcanomapToMaposcanAndMap ::+  ( MonadFreshNames m,+    Buildable rep,+    Op rep ~ SOAC rep+  ) =>+  Pat (LetDec rep) ->+  ( SubExp,+    Lambda rep,+    [Scan rep],+    Lambda rep,+    [VName]+  ) ->+  m (Stm rep, Stm rep)+maposcanomapToMaposcanAndMap (Pat pes) (w, post_lam, scans, map_lam, arrs) = do+  map_res <- mapM tempRes $ lambdaReturnType map_lam+  map_stm <- mkLet map_res . Op . Screma w arrs <$> mapSOAC map_lam+  id_lam <- mkIdentityLambda $ lambdaReturnType map_lam+  let res = patElemIdent <$> pes+      pre_arrs = map identName map_res+  post_stm <- mkLet res . Op . Screma w pre_arrs <$> maposcanomapSOAC id_lam scans post_lam+  pure (map_stm, post_stm)+  where+    tempRes res = newIdent "temp_res" $ res `arrayOfRow` w+ -- | Turn a Screma into a maposcanomap (possibly with mapout parts) and a -- Redomap.  This is used to handle Scremas that are so complicated -- that we cannot directly generate efficient parallel code for them.@@ -335,6 +361,52 @@       =<< mapSOAC map_lam    letTupExp desc $ WithAcc [(acc_shape, [v], Nothing) | v <- dest] withacc_lam++-- | Perform a histogram-like operation using accumulators and map.+doHist ::+  (MonadBuilder m, Buildable (Rep m), Op (Rep m) ~ SOAC (Rep m)) =>+  Name ->+  [HistOp (Rep m)] ->+  [VName] ->+  ([LParam (Rep m)] -> m [SubExp]) ->+  m [[VName]]+doHist desc ops arrs mk = do+  (inputs, cert_ps, acc_ts) <- unzip3 <$> mapM onOp ops+  acc_ps <- mapM (newParam "acc_p") acc_ts+  arrs_ts <- mapM lookupType arrs++  withacc_lam <- mkLambda (cert_ps <> acc_ps) $ do+    acc_ps_inner <- mapM (newParam "acc_p") acc_ts+    params <- mapM (newParam "v" . stripArray 1) arrs_ts+    map_lam <-+      mkLambda (acc_ps_inner <> params) $ do+        let num_is_per_op = map (shapeRank . histShape) ops+            num_vs_per_op = map (length . histDest) ops+        (is, vs) <- splitAt (sum num_is_per_op) <$> mk params+        let is_per_op = chunks num_is_per_op is+            vs_per_op = chunks num_vs_per_op vs+        fmap subExpsRes $ forM (zip3 acc_ps_inner is_per_op vs_per_op) $ \(acc_p_inner, op_is, op_vs) ->+          letSubExp "scatter_acc" . BasicOp $+            UpdateAcc Safe (paramName acc_p_inner) op_is op_vs++    let w = arraysSize 0 arrs_ts+    (fmap varsRes . letTupExp "acc_res")+      . Op+      . Screma w (map paramName acc_ps <> arrs)+      =<< mapSOAC map_lam+  fmap (chunks (map (length . histDest) ops)) $+    letTupExp desc $+      WithAcc inputs withacc_lam+  where+    onOp op = do+      idx_params <- replicateM (shapeRank (histShape op)) $ newParam "idx" $ Prim int64+      let addIdxParams lam = lam {lambdaParams = idx_params <> lambdaParams lam}+          input = (histShape op, histDest op, Just (addIdxParams (histOp op), histNeutral op))+          shape = histShape op+      elem_ts <- fmap (map (stripArray (shapeRank shape))) $ mapM lookupType $ histDest op+      cert_p <- newParam "acc_cert" $ Prim Unit+      let cert = paramName cert_p+      pure (input, cert_p, Acc cert shape elem_ts NoUniqueness)  -- | The most addition-like binary operator for some primitive type. addBinOp :: PrimType -> BinOp
src/Futhark/Transform/FirstOrderTransform.hs view
@@ -14,13 +14,15 @@     transformLambda,     transformSOAC,     transformScrema,+    transformFlatMap,   ) where  import Control.Monad import Control.Monad.State-import Data.List (find, zip4)+import Data.List (find, uncons, zip4) import Data.Map.Strict qualified as M+import Data.Maybe import Futhark.Analysis.Alias qualified as Alias import Futhark.IR qualified as AST import Futhark.IR.Prop.Aliases@@ -88,7 +90,11 @@ transformStmRecursively (Let pat aux (Op soac)) =   auxing aux $ transformSOAC pat =<< mapSOACM soacTransform soac   where-    soacTransform = identitySOACMapper {mapOnSOACLambda = transformLambda}+    soacTransform =+      identitySOACMapper+        { mapOnSOACLambda = transformLambda,+          mapOnSOACExtLambda = transformLambda+        } transformStmRecursively (Let pat aux e) =   auxing aux $ letBind pat =<< mapExpM transform e   where@@ -114,6 +120,194 @@         letExp "result" =<< eBlank t   mapM oneArray ts +-- | Sequentialise a single FlatMap. The size of the nonuniform results produced+-- by the lambda is not known until it has been run, so each of them is+-- accumulated in a scratch buffer that is doubled whenever it runs out of+-- space, and finally truncated to the actual size. The value results need no+-- such treatment, as there is exactly one per iteration. The shape and offset+-- arrays are filled in as we go, and the flag array is then a scatter of the+-- segment starts.+transformFlatMap ::+  (Transformer m) =>+  Pat (LetDec (Rep m)) ->+  SubExp ->+  [VName] ->+  ExtLambda (Rep m) ->+  m ()+transformFlatMap pat w arrs lam = do+  let irreg_ts = flatMapRowTypes lam+      reg_ts = flatMapUniformTypes lam+  arrs_ts <- mapM lookupType arrs++  -- Loop parameters: the current filled size, the current capacity, the+  -- per-element shape and offset arrays, one scratch buffer per nonuniform+  -- result, and one array per uniform result.+  size_p <- newParam "flatmap_size" $ toDecl (Prim int64) Nonunique+  cap_p <- newParam "flatmap_cap" $ toDecl (Prim int64) Nonunique+  shape_p <- newParam "flatmap_shape" $ toDecl (arrayOfRow (Prim int64) w) Unique+  offset_p <- newParam "flatmap_offset" $ toDecl (arrayOfRow (Prim int64) w) Unique+  scratch_ps <-+    forM irreg_ts $ \et ->+      newParam "flatmap_res" $ toDecl (arrayOfRow et (Var (paramName cap_p))) Unique+  reg_ps <-+    forM reg_ts $ \rt ->+      newParam "flatmap_reg" $ toDecl (arrayOfRow rt w) Unique++  -- The capacity initially matches the input size.+  shape_init <- letExp "flatmap_shape" $ BasicOp $ Scratch int64 [w]+  offset_init <- letExp "flatmap_offset" $ BasicOp $ Scratch int64 [w]+  scratch_init <- forM irreg_ts $ \et -> letExp "flatmap_res" =<< eBlank (arrayOfRow et w)+  reg_init <- forM reg_ts $ \rt -> letExp "flatmap_reg" =<< eBlank (arrayOfRow rt w)++  let merge =+        (size_p, intConst Int64 0)+          : (cap_p, w)+          : (shape_p, Var shape_init)+          : (offset_p, Var offset_init)+          : zip scratch_ps (map Var scratch_init)+            <> zip reg_ps (map Var reg_init)+      merge_params = map fst merge++  i <- newVName "i"+  let loop_form = ForLoop i Int64 w+      loop_scope = scopeOfLoopForm loop_form <> scopeOfFParams merge_params+  loop_body <- runBodyBuilder . localScope loop_scope $ do+    -- Apply the lambda to the current elements.+    let arg arr arr_t = BasicOp $ Index arr $ fullSlice arr_t [DimFix $ Var i]+        size = Var $ paramName size_p+    lam_res <- map resSubExp <$> bindLambda lam (zipWith arg arrs arrs_ts)++    -- The lambda produces the common length of its segment results first.+    let (k, ys) =+          fromMaybe (error "transformFlatMap: malformed FlatMap.") $+            uncons lam_res+        (irreg_ys, reg_ys) = flatMapSplitValues lam ys+    new_size <-+      letSubExp "flatmap_new_size" . BasicOp $+        BinOp (Add Int64 OverflowUndef) size k+    grow <-+      letSubExp "flatmap_grow" . BasicOp $+        CmpOp (CmpSlt Int64) (Var (paramName cap_p)) new_size++    -- New capacity: double it (but at least fit) when it overflows.+    new_cap <-+      letSubExp "flatmap_new_cap"+        =<< eIf+          (eSubExp grow)+          ( buildBody_ $ do+              doubled <-+                letSubExp "doubled" . BasicOp $+                  BinOp (Mul Int64 OverflowUndef) (Var (paramName cap_p)) (intConst Int64 2)+              fmap (pure . subExpRes) . letSubExp "atleast" . BasicOp $+                BinOp (SMax Int64) doubled new_size+          )+          (buildBody_ $ pure [subExpRes $ Var (paramName cap_p)])++    let lowSlice t =+          fullSlice t [DimSlice (intConst Int64 0) size (intConst Int64 1)]++    -- Grow (and copy) each scratch buffer when necessary, then write the+    -- new elements at the end.+    scratch_res <- forM (zip3 scratch_ps irreg_ts irreg_ys) $ \(sp, et, ys_j) -> do+      let full_t = arrayOfRow et new_cap+      base <-+        letExp "flatmap_grown"+          =<< eIf+            (eSubExp grow)+            ( buildBody_ $ do+                fresh <-+                  letExp "flatmap_fresh" . BasicOp $+                    Scratch (elemType full_t) (arrayDims full_t)+                old_t <- lookupType $ paramName sp+                copied <-+                  letInPlace "flatmap_fresh" fresh (lowSlice full_t) $+                    BasicOp (Index (paramName sp) (lowSlice old_t))+                pure [varRes copied]+            )+            ( buildBody_ $+                fmap (pure . varRes) . letExp "flatmap_kept" $+                  shapeCoerce (arrayDims full_t) (paramName sp)+            )+      letInPlace "flatmap_res" base (fullSlice full_t [DimSlice size k (intConst Int64 1)]) $+        BasicOp (SubExp ys_j)++    -- The segment's size and its offset (the running total before it).+    shape' <-+      letInPlace "flatmap_shape" (paramName shape_p) (fullSlice (paramType shape_p) [DimFix $ Var i]) $+        BasicOp (SubExp k)+    offset' <-+      letInPlace "flatmap_offset" (paramName offset_p) (fullSlice (paramType offset_p) [DimFix $ Var i]) $+        BasicOp (SubExp size)++    -- The uniform results are simply written at this iteration's index.+    reg_res <- forM (zip reg_ps reg_ys) $ \(rp, reg_y) ->+      letInPlace "flatmap_reg" (paramName rp) (fullSlice (paramType rp) [DimFix $ Var i]) $+        BasicOp (SubExp reg_y)++    pure $+      subExpsRes [new_size, new_cap]+        <> varsRes (shape' : offset' : scratch_res <> reg_res)++  loop_res <- letTupExp "flatmap" $ Loop merge loop_form loop_body+  case (loop_res, patNames pat) of+    (size_res : _cap_res : shape_res : offset_res : value_res, m_pat : shape_pat : flag_pat : offset_pat : out_pats) -> do+      -- Bind the total size and the shape/offset arrays, then truncate each+      -- buffer. The uniform results are already of the right size.+      letBindNames [m_pat] $ BasicOp $ SubExp $ Var size_res+      letBindNames [shape_pat] $ BasicOp $ SubExp $ Var shape_res+      letBindNames [offset_pat] $ BasicOp $ SubExp $ Var offset_res+      let (scratch_res, reg_res) = splitAt (length irreg_ts) value_res+          (data_pats, reg_pats) = flatMapSplitValues lam out_pats+      forM_ (zip data_pats scratch_res) $ \(out, scratch) -> do+        scratch_t <- lookupType scratch+        letBindNames [out] . BasicOp . Index scratch $+          fullSlice scratch_t [DimSlice (intConst Int64 0) (Var m_pat) (intConst Int64 1)]+      forM_ (zip reg_pats reg_res) $ \(out, reg) ->+        letBindNames [out] $ BasicOp $ SubExp $ Var reg+      -- The flag array: scatter a 'true' at the offset of each non-empty+      -- segment, over an otherwise 'false' array.+      transformFlatMapFlags flag_pat w (Var m_pat) shape_res offset_res+    _ ->+      error "transformFlatMap: malformed FlatMap."++-- | Compute a 'FlatMap' flag array of length @m@: 'true' at the start of each+-- non-empty segment, 'false' elsewhere. Emitted as a sequential scatter loop.+transformFlatMapFlags ::+  (Transformer m) =>+  VName ->+  SubExp ->+  SubExp ->+  VName ->+  VName ->+  m ()+transformFlatMapFlags flag_pat w m shape offset = do+  let flag_t = arrayOfRow (Prim Bool) m+  flags_init <- letExp "flatmap_flags" $ BasicOp $ Replicate (Shape [m]) (constant False)+  flags_p <- newParam "flatmap_flags" $ toDecl flag_t Unique+  j <- newVName "j"+  let flag_form = ForLoop j Int64 w+  shape_t <- lookupType shape+  offset_t <- lookupType offset+  flag_body <- runBodyBuilder $+    localScope (scopeOfLoopForm flag_form <> scopeOfFParams [flags_p]) $ do+      sz <- letSubExp "flatmap_sz" $ BasicOp $ Index shape $ fullSlice shape_t [DimFix $ Var j]+      off <- letSubExp "flatmap_off" $ BasicOp $ Index offset $ fullSlice offset_t [DimFix $ Var j]+      nonempty <-+        letSubExp "flatmap_nonempty" . BasicOp $+          CmpOp (CmpSlt Int64) (intConst Int64 0) sz+      flags' <-+        letSubExp "flatmap_flags"+          =<< eIf+            (eSubExp nonempty)+            ( buildBody_+                $ fmap (pure . varRes)+                  . letInPlace "flatmap_flags" (paramName flags_p) (fullSlice flag_t [DimFix off])+                $ BasicOp (SubExp (constant True))+            )+            (buildBody_ $ pure [varRes $ paramName flags_p])+      pure [subExpRes flags']+  letBindNames [flag_pat] $ Loop [(flags_p, Var flags_init)] flag_form flag_body+ -- | Sequentialise a single Screma. transformScrema ::   (Transformer m) =>@@ -237,6 +431,8 @@   error "transformSOAC: unhandled VJP" transformSOAC _ WithVJP {} =   error "transformSOAC: unhandled WithVJP"+transformSOAC pat (FlatMap w arrs lam) =+  transformFlatMap pat w arrs lam transformSOAC pat (Screma w arrs form) =   transformScrema pat w arrs form transformSOAC pat (Stream w arrs nes lam) = do@@ -373,8 +569,8 @@     LetDec rep ~ LetDec SOACS,     Alias.AliasableRep rep   ) =>-  Lambda SOACS ->-  m (AST.Lambda rep)+  GLambda SOACS t ->+  m (AST.GLambda rep t) transformLambda (Lambda params rettype body) = do   body' <-     fmap fst . runBuilder $@@ -395,7 +591,7 @@  bindLambda ::   (Transformer m) =>-  AST.Lambda (Rep m) ->+  AST.GLambda (Rep m) t ->   [AST.Exp (Rep m)] ->   m Result bindLambda (Lambda params _ body) args = do
+ src/Futhark/Transform/ISRWIM.hs view
@@ -0,0 +1,195 @@+{-# LANGUAGE TypeFamilies #-}++-- | Interchanging scans with inner maps.+module Futhark.Transform.ISRWIM+  ( iswim,+    irwim,+    rwimPossible,+  )+where++import Control.Monad+import Futhark.IR.SOACS+import Futhark.MonadFreshNames+import Futhark.Tools++-- | Interchange Scan With Inner Map. Tries to turn a @scan(map)@ into a+-- @map(scan)+iswim ::+  (MonadBuilder m, Rep m ~ SOACS) =>+  Pat Type ->+  SubExp ->+  Lambda SOACS ->+  [(SubExp, VName)] ->+  Maybe (m ())+iswim res_pat w scan_fun scan_input+  | Just (map_pat, map_aux, map_w, map_fun) <- rwimPossible scan_fun = Just $ do+      let (accs, arrs) = unzip scan_input+      let indexAcc (Var v) = do+            v_t <- lookupType v+            letSubExp "acc" $+              BasicOp $+                Index v $+                  fullSlice v_t [DimFix $ intConst Int64 0]+          indexAcc Constant {} =+            error "irwim: array accumulator is a constant."+      arrs' <- transposedArrays arrs+      accs' <- mapM indexAcc accs++      -- let (_red_acc_params, red_elem_params) =+      --       splitAt (length arrs) $ lambdaParams red_fun+      --     map_rettype = map rowType $ lambdaReturnType red_fun+      --     map_params = map (setParamOuterDimTo w) red_elem_params++      let map_arrs' = arrs'+          (_scan_acc_params, scan_elem_params) =+            splitAt (length arrs) $ lambdaParams scan_fun+          map_params = map (setParamOuterDimTo w) scan_elem_params+          map_rettype = map (setOuterDimTo w) $ lambdaReturnType scan_fun++          scan_params = lambdaParams map_fun+          scan_body = lambdaBody map_fun+          scan_rettype = lambdaReturnType map_fun+          scan_fun' = Lambda scan_params scan_rettype scan_body+          scan_input' = zip accs' $ map paramName map_params++      scan_soac <- scanSOAC [Scan scan_fun' accs']+      let map_body =+            mkBody+              ( oneStm $+                  Let (setPatOuterDimTo w map_pat) (defAux ()) $+                    Op $+                      Screma w (map snd scan_input') scan_soac+              )+              $ varsRes+              $ patNames map_pat+          map_fun' = Lambda map_params map_rettype map_body++      res_pat' <-+        fmap basicPat $+          mapM (newIdent' (<> "_transposed") . transposeIdentType) $+            patIdents res_pat++      addStm . Let res_pat' map_aux . Op . Screma map_w map_arrs'+        =<< mapSOAC map_fun'++      forM_ (zip (patIdents res_pat) (patIdents res_pat')) $ \(to, from) -> do+        let perm = [1, 0] ++ [2 .. arrayRank (identType from) - 1]+        addStm $+          Let (basicPat [to]) (defAux ()) . BasicOp $+            Rearrange (identName from) perm+  | otherwise = Nothing++-- | Interchange Reduce With Inner Map. Tries to turn a @reduce(map)@ into a+-- @map(reduce)+irwim ::+  (MonadBuilder m, Rep m ~ SOACS) =>+  Pat Type ->+  SubExp ->+  Commutativity ->+  Lambda SOACS ->+  [(SubExp, VName)] ->+  Maybe (m ())+irwim res_pat w comm red_fun red_input+  | Just (map_pat, map_aux, map_w, map_fun) <- rwimPossible red_fun = Just $ do+      let (accs, arrs) = unzip red_input+      arrs' <- transposedArrays arrs+      -- FIXME?  Can we reasonably assume that the accumulator is a+      -- replicate?  We also assume that it is non-empty.+      let indexAcc (Var v) = do+            v_t <- lookupType v+            letSubExp "acc" $+              BasicOp $+                Index v $+                  fullSlice v_t [DimFix $ intConst Int64 0]+          indexAcc Constant {} =+            error "irwim: array accumulator is a constant."+      accs' <- mapM indexAcc accs++      let (_red_acc_params, red_elem_params) =+            splitAt (length arrs) $ lambdaParams red_fun+          map_rettype = map rowType $ lambdaReturnType red_fun+          map_params = map (setParamOuterDimTo w) red_elem_params++          red_params = lambdaParams map_fun+          red_body = lambdaBody map_fun+          red_rettype = lambdaReturnType map_fun+          red_fun' = Lambda red_params red_rettype red_body+          red_input' = zip accs' $ map paramName map_params+          red_pat = stripPatOuterDim map_pat++      map_body <-+        case irwim red_pat w comm red_fun' red_input' of+          Nothing -> do+            reduce_soac <- reduceSOAC [Reduce comm red_fun' $ map fst red_input']+            pure+              $ mkBody+                ( oneStm $+                    Let red_pat (defAux ()) $+                      Op $+                        Screma w (map snd red_input') reduce_soac+                )+              $ varsRes+              $ patNames map_pat+          Just m -> localScope (scopeOfLParams map_params) $ do+            map_body_stms <- collectStms_ m+            pure $ mkBody map_body_stms $ varsRes $ patNames map_pat++      let map_fun' = Lambda map_params map_rettype map_body++      addStm . Let res_pat map_aux . Op . Screma map_w arrs'+        =<< mapSOAC map_fun'+  | otherwise = Nothing++-- | Does this reduce operator contain an inner map, and if so, what+-- does that map look like?+rwimPossible ::+  Lambda SOACS ->+  Maybe (Pat Type, StmAux (), SubExp, Lambda SOACS)+rwimPossible fun+  | Body _ stms res <- lambdaBody fun,+    [stm] <- stmsToList stms, -- Body has a single binding+    map_pat <- stmPat stm,+    map Var (patNames map_pat) == map resSubExp res, -- Returned verbatim+    Op (Screma map_w map_arrs form) <- stmExp stm,+    Just map_fun <- isMapSOAC form,+    map paramName (lambdaParams fun) == map_arrs =+      Just (map_pat, stmAux stm, map_w, map_fun)+  | otherwise =+      Nothing++transposedArrays :: (MonadBuilder m) => [VName] -> m [VName]+transposedArrays arrs = forM arrs $ \arr -> do+  t <- lookupType arr+  let perm = [1, 0] ++ [2 .. arrayRank t - 1]+  letExp (baseName arr) $ BasicOp $ Rearrange arr perm++setParamOuterDimTo :: SubExp -> LParam SOACS -> LParam SOACS+setParamOuterDimTo w param =+  let t = setOuterDimTo w $ paramType param+   in param {paramDec = t}++setIdentOuterDimTo :: SubExp -> Ident -> Ident+setIdentOuterDimTo w ident =+  let t = setOuterDimTo w $ identType ident+   in ident {identType = t}++setOuterDimTo :: SubExp -> Type -> Type+setOuterDimTo w t =+  arrayOfRow (rowType t) w++setPatOuterDimTo :: SubExp -> Pat Type -> Pat Type+setPatOuterDimTo w pat =+  basicPat $ map (setIdentOuterDimTo w) $ patIdents pat++transposeIdentType :: Ident -> Ident+transposeIdentType ident =+  ident {identType = transposeType $ identType ident}++stripIdentOuterDim :: Ident -> Ident+stripIdentOuterDim ident =+  ident {identType = rowType $ identType ident}++stripPatOuterDim :: Pat Type -> Pat Type+stripPatOuterDim pat =+  basicPat $ map stripIdentOuterDim $ patIdents pat
src/Futhark/Transform/Rename.hs view
@@ -96,9 +96,9 @@ -- correct to begin with.  Any free variables are left untouched. -- Note in particular that the parameters of the lambda are renamed. renameLambda ::-  (Renameable rep, MonadFreshNames m) =>-  Lambda rep ->-  m (Lambda rep)+  (Renameable rep, MonadFreshNames m, Rename t) =>+  GLambda rep t ->+  m (GLambda rep t) renameLambda = modifyNameSource . runRenamer . rename  -- | Produce an equivalent pattern but with each pattern element given@@ -311,7 +311,7 @@   rename (Acc acc ispace ts u) =     Acc <$> rename acc <*> rename ispace <*> rename ts <*> pure u -instance (Renameable rep) => Rename (Lambda rep) where+instance (Renameable rep, Rename t) => Rename (GLambda rep t) where   rename (Lambda params ret body) =     renameBound (map paramName params) $       Lambda <$> mapM rename params <*> mapM rename ret <*> rename body
src/Futhark/Transform/Substitute.hs view
@@ -172,7 +172,7 @@   substituteNames _ (Mem space) =     Mem space -instance (Substitutable rep) => Substitute (Lambda rep) where+instance (Substitutable rep, Substitute t) => Substitute (GLambda rep t) where   substituteNames substs (Lambda params rettype body) =     Lambda       (substituteNames substs params)
+ src/Futhark/Transform/ToGPU.hs view
@@ -0,0 +1,87 @@+module Futhark.Transform.ToGPU+  ( getSize,+    segThread,+    soacsLambdaToGPU,+    soacsStmToGPU,+    soacsExpToGPU,+    scopeForGPU,+    scopeForSOACs,+    injectSOACS,+  )+where++import Control.Monad.Identity+import Data.List ()+import Futhark.IR+import Futhark.IR.GPU+import Futhark.IR.SOACS (SOACS)+import Futhark.IR.SOACS.SOAC qualified as SOAC+import Futhark.Tools++getSize ::+  (MonadBuilder m, Op (Rep m) ~ HostOp inner (Rep m)) =>+  Name ->+  SizeClass ->+  m SubExp+getSize desc size_class = do+  size_key <- nameFromText . prettyText <$> newVName desc+  letSubExp desc $ Op $ SizeOp $ GetSize size_key size_class++segThread ::+  (MonadBuilder m, Rep m ~ GPU) =>+  Name ->+  m SegLevel+segThread desc =+  SegThread SegVirt <$> (Just <$> kernelGrid)+  where+    kernelGrid =+      KernelGrid+        <$> (Count <$> getSize (desc <> "_num_tblocks") SizeGrid)+        <*> (Count <$> getSize (desc <> "_tblock_size") SizeThreadBlock)++injectSOACS ::+  ( Monad m,+    SameScope from to,+    ExpDec from ~ ExpDec to,+    BodyDec from ~ BodyDec to,+    RetType from ~ RetType to,+    BranchType from ~ BranchType to,+    Op from ~ SOAC from+  ) =>+  (SOAC to -> Op to) ->+  Rephraser m from to+injectSOACS f =+  Rephraser+    { rephraseExpDec = pure,+      rephraseBodyDec = pure,+      rephraseLetBoundDec = pure,+      rephraseFParamDec = pure,+      rephraseLParamDec = pure,+      rephraseOp = fmap f . onSOAC,+      rephraseRetType = pure,+      rephraseBranchType = pure+    }+  where+    onSOAC = SOAC.mapSOACM mapper+    mapper =+      SOAC.SOACMapper+        { SOAC.mapOnSOACSubExp = pure,+          SOAC.mapOnSOACVName = pure,+          SOAC.mapOnSOACLambda = rephraseLambda $ injectSOACS f,+          SOAC.mapOnSOACExtLambda = rephraseLambda $ injectSOACS f+        }++soacsStmToGPU :: Stm SOACS -> Stm GPU+soacsStmToGPU = runIdentity . rephraseStm (injectSOACS OtherOp)++soacsExpToGPU :: Exp SOACS -> Exp GPU+soacsExpToGPU = runIdentity . rephraseExp (injectSOACS OtherOp)++soacsLambdaToGPU :: Lambda SOACS -> Lambda GPU+soacsLambdaToGPU = runIdentity . rephraseLambda (injectSOACS OtherOp)++scopeForSOACs :: Scope GPU -> Scope SOACS+scopeForSOACs = castScope++scopeForGPU :: Scope SOACS -> Scope GPU+scopeForGPU = castScope
src/Futhark/Util.hs view
@@ -16,12 +16,14 @@     chunks,     chunkLike,     dropAt,+    without,     takeLast,     dropLast,     mapEither,     partitionMaybe,     maybeNth,     maybeHead,+    unsnoc,     lookupWithIndex,     splitFromEnd,     splitAt3,@@ -165,6 +167,10 @@ dropAt :: Int -> Int -> [a] -> [a] dropAt i n xs = take i xs ++ drop (i + n) xs +-- | Remove the element at the given index.+without :: Int -> [a] -> [a]+without j = dropAt j 1+ -- | @takeLast n l@ takes the last @n@ elements of @l@. takeLast :: Int -> [a] -> [a] takeLast n = reverse . take n . reverse@@ -197,6 +203,12 @@ maybeHead :: [a] -> Maybe a maybeHead [] = Nothing maybeHead (x : _) = Just x++-- | Split the last element from the list, if it exists.+unsnoc :: [a] -> Maybe ([a], a)+unsnoc [] = Nothing+unsnoc [x] = Just ([], x)+unsnoc (x : xs) = unsnoc xs >>= \(ys, y) -> Just (x : ys, y)  -- | Lookup a value, returning also the index at which it appears. lookupWithIndex :: (Eq a) => a -> [(a, b)] -> Maybe (Int, b)
src/Language/Futhark/Interpreter.hs view
@@ -746,6 +746,39 @@ typeValueShape :: Env -> StructType -> EvalM ValueShape typeValueShape env t = typeShape <$> evalTypeFully (expandType env t) +-- | Compute the shape of a sum value with the given type, constructor+-- and payload values. The shapes of the payload values take+-- precedence over the type annotation, which may contain existential+-- sizes that cannot be evaluated. The shapes of the other+-- constructors are computed from the type, after resolving+-- existential sizes against the concrete payload where possible; the+-- type system guarantees that dimensions not determined by the+-- payload can be evaluated.+sumValueShape :: Env -> StructType -> Name -> [Value] -> EvalM ValueShape+sumValueShape env t c vs =+  case expandType env t of+    Scalar (Sum cs) -> do+      let payload_shapes = map valueShape vs+          learned =+            i64Env $ mconcat $ maybe [] (zipWith learn payload_shapes) $ M.lookup c cs+          evalDim (SizeClosure denv e) =+            asInt64 <$> evalWithExts (learned <> denv) e+          onConstr c' fts+            | c' == c = pure payload_shapes+            | otherwise = map typeShape <$> mapM (bitraverse evalDim pure) fts+      ShapeSum <$> M.traverseWithKey onConstr cs+    t' -> typeShape <$> evalTypeFully t'+  where+    learn (ShapeDim d s) ft+      | SizeClosure _ (Var v _ _) : _ <- shapeDims (arrayShape ft) =+          M.insert (qualLeaf v) d $ learn s (stripArray 1 ft)+      | otherwise = learn s (stripArray 1 ft)+    learn (ShapeRecord fs) (Scalar (Record fts)) =+      mconcat $ M.elems $ M.intersectionWith learn fs fts+    learn (ShapeSum fs) (Scalar (Sum fts)) =+      mconcat $ map mconcat $ M.elems $ M.intersectionWith (zipWith learn) fs fts+    learn _ _ = mempty+ -- Sometimes type instantiation is not quite enough - then we connect -- up the missing sizes here.  In particular used for eta-expanded -- entry points.@@ -772,7 +805,7 @@         env'' <- linkMissingSizes missing_sizes p v <$> matchPat env' p v         etaExpand (v : vs) env'' rt     etaExpand vs env' _ = do-      f <- eval env' body+      f <- localExts $ eval env' body       foldM (apply noLoc mempty) f $ reverse vs evalBinding env missing_sizes (p : ps) body rettype =   pure . ValueFun $ \v -> do@@ -781,37 +814,59 @@  evalValBinding ::   Env ->+  VName ->   [TypeParam] ->   [Pat ParamType] ->   ResRetType ->   Exp ->   EvalM TermBinding-evalValBinding env tparams ps ret fbody = do+evalValBinding env name tparams ps ret fbody = do   let ftype = evalToStruct $ expandType env $ funType ps ret       retext = case ps of         [] -> retDims ret         _ -> [] -  -- Distinguish polymorphic and non-polymorphic bindings here.-  if null tparams-    then+  -- A function binding may refer to itself recursively, which means it muse be+  -- in scope of itself. We tie the knot by evaluating the body in an+  -- environment ('recenv') that maps the 'name' to itself. Crucially, in each+  -- case below the binding is a lazy, self-referential 'let' whose right-hand+  -- side is a value constructor (a 'ValueFun' closure or a 'TermPoly' closure).+  --+  -- Distinguish polymorphic, monomorphic-function, and plain-value+  -- bindings here.+  case (null tparams, ps) of+    (True, []) ->+      -- Not a syntactic function, so cannot be recursive.       fmap (TermValue (Just $ T.BoundV [] ftype))         . returned env (retType ret) retext         =<< evalBinding env [] ps fbody (retType ret)-    else pure . TermPoly (Just $ T.BoundV [] ftype) $ \ftype' -> do-      let resolved =-            resolveTypeParams (map typeParamName tparams) ftype ftype'-      tparam_env <- evalResolved resolved-      let env' = tparam_env <> env-          -- In some cases (abstract lifted types) there may be-          -- missing sizes that were not fixed by the type-          -- instantiation.  These will have to be set by looking-          -- at the actual function arguments.-          missing_sizes =-            filter (`M.notMember` envTerm env') $-              map typeParamName (filter isSizeParam tparams)-      returned env (retType ret) retext-        =<< evalBinding env' missing_sizes ps fbody (retType ret)+    (True, p : ps') ->+      -- A monomorphic function; "unfold" the first step of what evalBinding+      -- would do.+      let recenv = env {envTerm = M.insert name binding $ envTerm env}+          binding =+            TermValue (Just $ T.BoundV [] ftype) . ValueFun $ \v -> do+              env' <- matchPat recenv p v+              evalBinding env' [] ps' fbody (retType ret)+       in pure binding+    (False, _) ->+      -- A polymorphic function.+      let binding = TermPoly (Just $ T.BoundV [] ftype) $ \ftype' -> do+            let resolved =+                  resolveTypeParams (map typeParamName tparams) ftype ftype'+            tparam_env <- evalResolved resolved+            let recenv = env {envTerm = M.insert name binding $ envTerm env}+                env' = tparam_env <> recenv+                -- In some cases (abstract lifted types) there may be+                -- missing sizes that were not fixed by the type+                -- instantiation.  These will have to be set by looking+                -- at the actual function arguments.+                missing_sizes =+                  filter (`M.notMember` envTerm env') $+                    map typeParamName (filter isSizeParam tparams)+            returned env (retType ret) retext+              =<< evalBinding env' missing_sizes ps fbody (retType ret)+       in pure binding  evalArg :: Env -> Exp -> Maybe VName -> EvalM Value evalArg env e ext = do@@ -887,9 +942,9 @@       env'' = env' <> i64Env (resolveExistentials (map sizeName sizes) p_t v_s)   eval env'' body evalAppExp env (LetFun (f, _) (tparams, ps, _, Info ret, fbody) body _) = do-  binding <- evalValBinding env tparams ps ret fbody+  binding <- evalValBinding env f tparams ps ret fbody   eval (env {envTerm = M.insert f binding $ envTerm env}) body-evalAppExp env (BinOp (op, _) op_t (x, Info xext) (y, Info yext) loc)+evalAppExp env (BinOp (op, _) (Info op_t) (x, Info xext) (y, Info yext) loc)   | baseName (qualLeaf op) == "&&" = do       x' <- asBool <$> eval env x       if x'@@ -903,8 +958,9 @@   | otherwise = do       x' <- evalArg env x xext       y' <- evalArg env y yext-      op' <- eval env $ Var op op_t loc-      apply2 loc env op' x' y'+      op' <- evalTermVar env op op_t+      op'' <- apply loc env op' x'+      apply loc env op'' y' evalAppExp env (If cond e1 e2 _) = do   cond' <- asBool <$> eval env cond   if cond' then eval env e1 else eval env e2@@ -1078,14 +1134,19 @@   evalTermVar env qv $ toStruct t eval env (OpSectionLeft qv _ e (Info (_, _, argext), _) (Info (RetType _ t), _) loc) = do   v <- evalArg env e argext-  f <- evalTermVar env qv (toStruct t)+  f <- evalTermVar env qv t'   apply loc env f v+  where+    t' = toStruct t eval env (OpSectionRight qv _ e (Info _, Info (_, _, argext)) (Info (RetType _ t)) loc) = do   y <- evalArg env e argext   pure $     ValueFun $ \x -> do-      f <- evalTermVar env qv $ toStruct t-      apply2 loc env f x y+      f <- evalTermVar env qv t'+      f' <- apply loc env f x+      apply loc env f' y+  where+    t' = toStruct t eval env (UpdateSection steps _ loc) =   pure $ ValueFun $ evalSection steps   where@@ -1106,7 +1167,7 @@   eval env e eval env (Constr c es (Info t) _) = do   vs <- mapM (eval env) es-  shape <- typeValueShape env $ toStruct t+  shape <- sumValueShape env (toStruct t) c vs   pure $ ValueSum shape c vs eval env (Attr (AttrAtom (AtomName "break") _) e loc) = do   break env (locOf loc)@@ -1243,7 +1304,7 @@  evalDec :: Env -> Dec -> EvalM Env evalDec env (ValDec (ValBind _ v _ _ (Info ret) tparams ps fbody _ _ _)) = localExts $ do-  binding <- evalValBinding env tparams ps ret fbody+  binding <- evalValBinding env v tparams ps ret fbody   sizes <- extEnv   pure $ mempty {envTerm = M.singleton v binding} <> sizes evalDec env (OpenDec me _) = do@@ -1827,6 +1888,40 @@               error $                 "Invalid arguments to map intrinsic:\n"                   ++ unlines [prettyString t, show f, show xs]+    def "flatmap" = Just $+      TermPoly Nothing $ \t ->+        pure $ ValueFun $ \f -> pure . ValueFun $ \xs ->+          case unfoldFunType t of+            ([_, _], ret_t)+              | Just [_, _, _, irreg_t, reg_t] <- isTupleRecord ret_t -> do+                  irreg_rowshape <- typeShape <$> evalTypeFully (stripArray 1 irreg_t)+                  reg_rowshape <- typeShape <$> evalTypeFully (stripArray 1 reg_t)+                  yss <-+                    mapM+                      (apply noLoc mempty f)+                      (snd $ fromArray xs)+                  -- Each application produces a segment, which is concatenated+                  -- with the others, and a value that is merely collected.+                  let (segs, regs) = unzip $ map (fromPair . fromTuple) yss+                      seg_sizes = map (genericLength . snd . fromArray) segs :: [Int64]+                      offsets = init $ scanl (+) 0 seg_sizes+                      flag s = if s == 0 then [] else True : replicate (fromIntegral s - 1) False+                      mkI64 = ValuePrim . SignedValue . Int64Value+                  pure $+                    toTuple+                      [ toArray' ShapeLeaf $ map mkI64 seg_sizes,+                        toArray' ShapeLeaf $ map (ValuePrim . BoolValue) $ concatMap flag seg_sizes,+                        toArray' ShapeLeaf $ map mkI64 offsets,+                        toArray' irreg_rowshape $ concatMap (snd . fromArray) segs,+                        toArray' reg_rowshape regs+                      ]+            _ ->+              error $+                "Invalid arguments to flatmap intrinsic:\n"+                  ++ unlines [show f, show xs]+      where+        fromPair (Just [x, y]) = (x, y)+        fromPair _ = error "flatmap: lambda did not return a pair"     def s | "reduce" `T.isPrefixOf` s = Just $       fun3 $ \f ne xs ->         foldM (apply2 noLoc mempty f) ne $ snd $ fromArray xs@@ -2257,7 +2352,7 @@           </> "Got input of types"           </> indent 2 (stack (map pretty args_ts))   where-    (param_ts, _) = unfoldFunType entry_t+    param_ts = map snd $ fst $ unfoldFunType entry_t     args_ts = map (valueStructType . valueType) args     expected       | null param_ts =
src/Language/Futhark/Interpreter/Values.hs view
@@ -7,7 +7,6 @@     Shape (..),     ValueShape,     typeShape,-    structTypeShape,      -- * Values     Value (..),@@ -93,12 +92,6 @@       typeShape t'   | otherwise =       ShapeLeaf--structTypeShape :: StructType -> Shape (Maybe Int64)-structTypeShape = fmap dim . typeShape-  where-    dim (IntLit x _ _) = Just $ fromIntegral x-    dim _ = Nothing  -- | A fully evaluated Futhark value. data Value m
src/Language/Futhark/Pretty.hs view
@@ -8,6 +8,7 @@     leadingOperator,     symbolName,     IsName (..),+    prettyNameText,     prettyNameString,     Annot (..),   )@@ -55,9 +56,13 @@   prettyName = pretty   toName = id +-- | Prettyprint name as text.+prettyNameText :: (IsName v) => v -> T.Text+prettyNameText = docText . prettyName+ -- | Prettyprint name as string.  Only use this for debugging. prettyNameString :: (IsName v) => v -> String-prettyNameString = T.unpack . docText . prettyName+prettyNameString = T.unpack . prettyNameText  -- | Class for type constructors that represent annotations.  Used in -- the prettyprinter to either print the original AST, or the computed@@ -153,7 +158,7 @@  prettyType :: (Pretty (Shape dim), Pretty u) => Int -> TypeBase dim u -> Doc a prettyType _ (Array u shape at) =-  pretty u <> pretty shape <> align (prettyScalarType 1 at)+  pretty u <> pretty shape <> align (prettyScalarType 2 at) prettyType p (Scalar t) =   prettyScalarType p t @@ -229,7 +234,8 @@ letBody body = "in" <+> align (pretty body)  prettyAppExp :: (IsName vn, Annot f) => Int -> AppExpBase f vn -> Doc a-prettyAppExp p (BinOp (bop, _) _ (x, _) (y, _) _) = prettyBinOp p bop x y+prettyAppExp p (BinOp (bop, _) _ (x, _) (y, _) _) =+  prettyBinOp p bop x y prettyAppExp _ (Match e cs _) = "match" <+> pretty e </> (stack . map pretty) (NE.toList cs) prettyAppExp _ (Loop sizeparams pat initexp form loopbody _) =   "loop"@@ -306,7 +312,10 @@ prettyAppExp p (Apply f args _) =   parensIf (p >= 10) $     prettyExp 0 f-      <+> hsep (map (prettyExp 10 . snd) $ NE.toList args)+      <+> hsep (map prettyArg $ NE.toList args)+  where+    prettyArg (_, e) =+      prettyExp 10 e  prettyLetLhsUpdate :: (IsName vn, Annot f) => [UpdateStep f vn] -> Doc a prettyLetLhsUpdate = mconcat . map pp
src/Language/Futhark/Prop.hs view
@@ -318,7 +318,9 @@ arrayOfWithAliases u shape2 (Array _ shape1 et) =   Array u (shape2 <> shape1) et arrayOfWithAliases u shape (Scalar t) =-  Array u shape (second (const mempty) t)+  if shapeRank shape == 0+    then Scalar t `setUniqueness` u+    else Array u shape (second (const mempty) t)  -- | @stripArray n t@ removes the @n@ outermost layers of the array. -- Essentially, it is the type of indexing an array of type @t@ with@@ -510,7 +512,7 @@ typeOf (AppExp _ (Info res)) = appResType res  -- | The type of a function with the given parameters and return type.-funType :: [Pat ParamType] -> ResRetType -> StructType+funType :: [Pat (TypeBase d Diet)] -> RetTypeBase d Uniqueness -> TypeBase d NoUniqueness funType params ret =   let RetType _ t = foldr (arrow . patternParam) ret params    in toStruct t@@ -520,7 +522,7 @@  -- | @foldFunType ts ret@ creates a function type ('Arrow') that takes -- @ts@ as parameters and returns @ret@.-foldFunType :: [ParamType] -> ResRetType -> StructType+foldFunType :: [TypeBase d Diet] -> RetTypeBase d Uniqueness -> TypeBase d NoUniqueness foldFunType ps ret =   let RetType _ t = foldr arrow ret ps    in toStruct t@@ -530,10 +532,10 @@  -- | 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 Diet], TypeBase dim NoUniqueness)-unfoldFunType (Scalar (Arrow _ _ d t1 (RetType _ t2))) =+unfoldFunType :: TypeBase dim as -> ([(PName, TypeBase dim Diet)], TypeBase dim NoUniqueness)+unfoldFunType (Scalar (Arrow _ p d t1 (RetType _ t2))) =   let (ps, r) = unfoldFunType t2-   in (second (const d) t1 : ps, r)+   in ((p, second (const d) t1) : ps, r) unfoldFunType t = ([], toStruct t)  -- | The type scheme of a value binding, comprising the type@@ -630,7 +632,7 @@  -- | When viewed as a function parameter, does this pattern correspond -- to a named parameter of some type?-patternParam :: Pat ParamType -> (PName, Diet, StructType)+patternParam :: Pat (TypeBase d Diet) -> (PName, Diet, TypeBase d NoUniqueness) patternParam (PatParens p _) =   patternParam p patternParam (PatAttr _ p _) =@@ -1013,6 +1015,25 @@                   ]                   $ RetType []                   $ Scalar (t_b Nonunique)+              ),+              ( "flatmap",+                IntrinsicPolyFun+                  [tp_a, tp_b, tp_c, sp_n]+                  [ Scalar $+                      Arrow mempty Unnamed Observe (Scalar (t_a NoUniqueness)) $+                        RetType [k] . Scalar . tupleRecord $+                          [ array_b Nonunique (shape [k]),+                            Scalar $ t_c Nonunique+                          ],+                    array_a Observe $ shape [n]+                  ]+                  $ RetType [m] . Scalar . tupleRecord+                  $ [ Array Unique (shape [n]) (Prim $ Signed Int64),+                      Array Unique (shape [m]) (Prim Bool),+                      Array Unique (shape [n]) (Prim $ Signed Int64),+                      array_b Unique $ shape [m],+                      array_c Unique $ shape [n]+                    ]               )             ]               ++@@ -1155,7 +1176,7 @@      intrinsicStart = 1 + baseTag (fst $ last primOp) -    [a, b, n, m, k, l, p, q] = zipWith VName (map nameFromText ["a", "b", "n", "m", "k", "l", "p", "q"]) [0 ..]+    [a, b, c, n, m, k, l, p, q] = zipWith VName (map nameFromText ["a", "b", "c", "n", "m", "k", "l", "p", "q"]) [0 ..]      t_a u = TypeVar u (qualName a) []     array_a u s = Array u s $ t_a mempty@@ -1164,6 +1185,10 @@     t_b u = TypeVar u (qualName b) []     array_b u s = Array u s $ t_b mempty     tp_b = TypeParamType Unlifted b mempty++    t_c u = TypeVar u (qualName c) []+    array_c u s = Array u s $ t_c mempty+    tp_c = TypeParamType Unlifted c mempty      [sp_n, sp_m, sp_k, sp_l, sp_p, sp_q] = map (`TypeParamDim` mempty) [n, m, k, l, p, q]
src/Language/Futhark/Traversals.hs view
@@ -74,7 +74,7 @@     Match <$> mapOnExp tv e <*> astMap tv cases <*> pure loc   astMap tv (Apply f args loc) = do     f' <- mapOnExp tv f-    args' <- traverse (traverse $ mapOnExp tv) args+    args' <- traverse onArg args     -- Safe to disregard return type because existentials cannot be     -- instantiated here, as the return is necessarily a function.     pure $ case f' of@@ -82,6 +82,8 @@         Apply f_inner (args_inner <> args') loc       _ ->         Apply f' args' loc+    where+      onArg (ext, e) = (ext,) <$> mapOnExp tv e   astMap tv (LetPat sizes pat e body loc) =     LetPat sizes <$> astMap tv pat <*> mapOnExp tv e <*> mapOnExp tv body <*> pure loc   astMap tv (LetFun name (tparams, params, ret, t, e) body loc) =@@ -105,13 +107,15 @@     where       mapStep (UpdateStepSlice slice) = UpdateStepSlice <$> mapM (astMap tv) slice       mapStep (UpdateStepField f) = pure $ UpdateStepField f-  astMap tv (BinOp (fname, fname_loc) t (x, xext) (y, yext) loc) =+  astMap tv (BinOp (fname, fname_loc) t x y loc) =     BinOp       <$> ((,) <$> mapOnName tv fname <*> pure fname_loc)       <*> traverse (mapOnStructType tv) t-      <*> ((,) <$> mapOnExp tv x <*> pure xext)-      <*> ((,) <$> mapOnExp tv y <*> pure yext)+      <*> onArg x+      <*> onArg y       <*> pure loc+    where+      onArg (e, ext) = (,ext) <$> mapOnExp tv e   astMap tv (Loop sparams mergepat loopinit form loopbody loc) =     Loop sparams       <$> astMap tv mergepat
src/Language/Futhark/TypeChecker.hs view
@@ -443,6 +443,7 @@   (MTy p_abs p_mod, psig_e') <- checkModTypeExp psig_e   bindSpaced1 Term pname loc $ \pname' -> do     let in_body_env = mempty {envModTable = M.singleton pname' p_mod}+    addTySet p_abs     localEnv in_body_env $       m (ModParam pname' psig_e' (Info $ map qualLeaf $ M.keys p_abs) loc) p_abs p_mod @@ -632,7 +633,8 @@       ([], EntryType t te)  -- | Check that a type is non-functional, looking up the liftedness of abstract--- types in the environment. This works because entry points cannot be polymorphic, so any remaining type names must be abstract.+-- types in the environment. This works because entry points cannot be+-- polymorphic, so any remaining type names must be abstract. orderZeroM :: TypeBase dim u -> TypeM Bool orderZeroM t = do   (orderZero t &&) . and <$> mapM isUnlifted (typeQualVars t)@@ -677,7 +679,7 @@   where     (RetType _ rettype_t) = rettype     (rettype_params, rettype') = unfoldFunType rettype_t-    param_ts = map patternType params ++ rettype_params+    param_ts = map patternType params ++ map snd rettype_params  checkValBind :: ValBindBase NoInfo Name -> TypeM (Env, ValBind) checkValBind vb = do@@ -695,11 +697,12 @@     checkFunDef (fname, maybe_tdecl, tparams, params, body, loc)    let entry' = Info (entryPoint doc params' maybe_tdecl' rettype) <$ entry+      vb' = ValBind entry' fname fname_loc maybe_tdecl' (Info rettype) tparams' params' body' doc attrs' loc+   case entry' of     Just _ -> checkEntryPoint loc tparams' params' rettype     _ -> pure () -  let vb' = ValBind entry' fname fname_loc maybe_tdecl' (Info rettype) tparams' params' body' doc attrs' loc   pure     ( mempty         { envVtable =
+ src/Language/Futhark/TypeChecker/Constraints.hs view
@@ -0,0 +1,120 @@+-- | Constraints produced (and solved) by the type checker.+module Language.Futhark.TypeChecker.Constraints+  ( Reason (..),+    CtType,+    CtTy (..),+    TyVarInfo (..),+    Level,+    TyVar,+    TyVars,+    TyParams,+  )+where++import Data.Bifunctor+import Data.Loc+import Data.Map qualified as M+import Futhark.Util.Pretty+import Language.Futhark++-- | The type representation used by the constraint solver.+type CtType d = TypeBase d NoUniqueness++-- | The reason for a type constraint. Used to generate type error+-- messages. The expected type is always the first one.+data Reason t+  = -- | No particular reason.+    Reason Loc+  | -- | Arising from pattern match.+    ReasonPatMatch Loc (PatBase NoInfo VName StructType) t+  | -- | Arising from explicit ascription.+    ReasonAscription Loc t t+  | ReasonRetType Loc t t+  | -- | Arising from checking a function argument. The types are the+    -- expected (parameter) and actual (argument) types.+    ReasonApply Loc (Maybe (QualName VName), Int) Exp t t+  | -- | Used when unifying a type with a function type in a function+    -- application. If this unification fails, it means the supposed+    -- function was not a function after all.+    ReasonApplySplit Loc (Maybe (QualName VName), Int) Exp t+  | ReasonBranches Loc t t+  deriving (Eq, Show, Functor, Foldable, Traversable)++instance Located (Reason t) where+  locOf (Reason l) = l+  locOf (ReasonPatMatch l _ _) = l+  locOf (ReasonAscription l _ _) = l+  locOf (ReasonRetType l _ _) = l+  locOf (ReasonApply l _ _ _ _) = l+  locOf (ReasonApplySplit l _ _ _) = l+  locOf (ReasonBranches l _ _) = l++-- | A type constraint.+data CtTy d = CtEq (Reason (CtType d)) (TypeBase d NoUniqueness) (TypeBase d NoUniqueness)+  deriving (Show)++instance Functor CtTy where+  fmap f (CtEq r x y) = CtEq (fmap (first f) r) (first f x) (first f y)++ctReason :: CtTy d -> Reason (CtType d)+ctReason (CtEq r _ _) = r++instance Located (CtTy d) where+  locOf = locOf . ctReason++instance Pretty (CtTy Size) where+  pretty (CtEq _ t1 t2) = pretty t1 <+> "~" <+> pretty t2++instance Pretty (CtTy ()) where+  pretty (CtEq _ t1 t2) = pretty t1 <+> "~" <+> pretty t2++-- | Information about a flexible type variable. Every type variable+-- is associated with a location, which is the original syntax element+-- that it is the type of.+data TyVarInfo d+  = -- | Can be substituted with anything.+    TyVarFree Loc Liftedness+  | -- | Can only be substituted with these primitive types.+    TyVarPrim Loc [PrimType]+  | -- | Must be a record with these fields.+    TyVarRecord Loc (M.Map Name (CtType d))+  | -- | Must be a sum type with these fields.+    TyVarSum Loc (M.Map Name [CtType d])+  deriving (Show, Eq)++instance Functor TyVarInfo where+  fmap _ (TyVarFree loc l) = TyVarFree loc l+  fmap _ (TyVarPrim loc ts) = TyVarPrim loc ts+  fmap f (TyVarRecord loc m) = TyVarRecord loc $ M.map (first f) m+  fmap f (TyVarSum loc m) = TyVarSum loc $ M.map (map (first f)) m++prettyTyVarInfo :: (Pretty (Shape d)) => TyVarInfo d -> Doc a+prettyTyVarInfo (TyVarFree _ l) = "free" <+> pretty l+prettyTyVarInfo (TyVarPrim _ pts) = "∈" <+> pretty pts+prettyTyVarInfo (TyVarRecord _ fs) = pretty $ Scalar $ Record fs+prettyTyVarInfo (TyVarSum _ cs) = pretty $ Scalar $ Sum cs++instance Pretty (TyVarInfo ()) where+  pretty = prettyTyVarInfo++instance Located (TyVarInfo d) where+  locOf (TyVarFree loc _) = loc+  locOf (TyVarPrim loc _) = loc+  locOf (TyVarRecord loc _) = loc+  locOf (TyVarSum loc _) = loc++-- | The name of a type variable.+type TyVar = VName++-- | The level at which a type variable is bound.  Higher means+-- deeper.  We can only unify a type variable at level @i@ with a type+-- @t@ if all type names that occur in @t@ are at most at level @i@.+type Level = Int++-- | If a VName is not in this map, it should be in the 'TyParams' -+-- the exception is abstract types, which are just missing (and+-- assumed to have smallest possible level).+type TyVars d = M.Map TyVar (Level, TyVarInfo d)++-- | Explicit type parameters.+type TyParams = M.Map TyVar (Level, Liftedness, Loc)
src/Language/Futhark/TypeChecker/Consumption.hs view
@@ -13,7 +13,7 @@  import Control.Monad import Control.Monad.Reader-import Control.Monad.State+import Control.Monad.State.Strict import Data.Bifoldable import Data.Bifunctor import Data.DList qualified as DL@@ -508,7 +508,7 @@   where     check seen als = do       when (any (`S.member` seen) als) $-        addError loc mempty . withIndexLink "self-aliases-arg" $+        addError loc mempty . withIndexLink "self-aliasing-arg" $           "Argument passed for consuming parameter is self-aliased."       pure $ als <> seen @@ -527,7 +527,7 @@   consumed e_cons   let e_t = typeOf e'   when (e_cons /= mempty && not (orderZero e_t)) $-    addError (locOf e) mempty $+    addError (locOf e) mempty . withIndexLink "consuming-argument" $       "Argument of functional type"         </> indent 2 (pretty e_t)         </> "contains consumption, which is not allowed."@@ -643,18 +643,22 @@   let checkMergeReturn (Id pat_v (Info pat_v_t) patloc) t = do         let free_als = S.filter (`notElem` patNames param) $ boundAliases (aliases t)         when (diet pat_v_t == Consume) $ forM_ free_als $ \v ->-          lift . addError loop_loc mempty $-            "Return value for consuming loop parameter"+          lift+            . addError loop_loc mempty+            . withIndexLink "consuming-loop-param-aliases"+            $ "Return value for consuming loop parameter"               <+> dquotes (prettyName pat_v)               <+> "aliases"               <+> dquotes (prettyName v)               <> "."         (cons, obs) <- get-        unless (S.null $ aliases t `S.intersection` cons) $-          lift . addError loop_loc mempty $-            "Return value for loop parameter"-              <+> dquotes (prettyName pat_v)-              <+> "aliases other consumed loop parameter."+        unless (S.null $ aliases t `S.intersection` cons)+          $ lift+            . addError loop_loc mempty+            . withIndexLink "loop-parameter-aliases-other"+          $ "Return value for loop parameter"+            <+> dquotes (prettyName pat_v)+            <+> "aliases other consumed loop parameter."         when           ( diet pat_v_t == Consume               && not (S.null (aliases t `S.intersection` (cons <> obs)))@@ -807,7 +811,7 @@   consumed e_cons   let e_t = typeOf e'   when (e_cons /= mempty && not (orderZero e_t)) $-    addError (locOf e) mempty $+    addError (locOf e) mempty . withIndexLink "contains-consumption" $       "Let-bound expression of higher-order type"         </> indent 2 (pretty e_t)         </> "contains consumption, which is not allowed."@@ -873,7 +877,7 @@ -- checkExp (AppExp (BinOp (op, oploc) opt (x, xp) (y, yp) loc) appres) = do   op_als <- observeVar (locOf oploc) (qualLeaf op) (unInfo opt)-  let at1 : at2 : _ = fst $ unfoldFunType op_als+  let (_, at1) : (_, at2) : _ = fst $ unfoldFunType op_als   (x', x_als) <- checkArg [] at1 x   (y', y_als) <- checkArg [(x', x_als)] at2 y   res_als <- checkFuncall loc (Just op) op_als [x_als, y_als]
+ src/Language/Futhark/TypeChecker/Error.hs view
@@ -0,0 +1,68 @@+-- | Fundamental facilities for constructing type error messages.+module Language.Futhark.TypeChecker.Error+  ( -- * Breadcrumbs+    BreadCrumbs,+    hasNoBreadCrumbs,+    matchingField,+    matchingConstructor,+    matching,+  )+where++import Futhark.Util.Pretty+import Language.Futhark++-- | A piece of information that describes what process the type+-- checker currently performing.  This is used to give better error+-- messages for unification errors.+data BreadCrumb+  = MatchingFields [Name]+  | MatchingConstructor Name+  | Matching (Doc ())++instance Pretty BreadCrumb where+  pretty (MatchingFields fields) =+    "When matching types of record field"+      <+> dquotes (mconcat $ punctuate "." $ map pretty fields)+      <> dot+  pretty (MatchingConstructor c) =+    "When matching types of constructor" <+> dquotes (pretty c) <> dot+  pretty (Matching s) =+    unAnnotate s++-- | Unification failures can occur deep down inside complicated types+-- (consider nested records). We leave breadcrumbs behind us so we can+-- report the path we took to find the mismatch. When combining+-- breadcrumbs with the 'Semigroup' instance, put the innermost+-- breadcrumbs to the left.+newtype BreadCrumbs = BreadCrumbs [BreadCrumb]++instance Semigroup BreadCrumbs where+  BreadCrumbs (MatchingFields xs : bcs1) <> BreadCrumbs (MatchingFields ys : bcs2) =+    BreadCrumbs $ MatchingFields (ys <> xs) : bcs1 <> bcs2+  BreadCrumbs bcs1 <> BreadCrumbs bcs2 =+    BreadCrumbs $ bcs1 <> bcs2++instance Monoid BreadCrumbs where+  mempty = BreadCrumbs []++-- | Is the path empty?+hasNoBreadCrumbs :: BreadCrumbs -> Bool+hasNoBreadCrumbs (BreadCrumbs []) = True+hasNoBreadCrumbs _ = False++-- | Matching a record field.+matchingField :: Name -> BreadCrumbs+matchingField f = BreadCrumbs [MatchingFields [f]]++-- | Matching a constructor.+matchingConstructor :: Name -> BreadCrumbs+matchingConstructor c = BreadCrumbs [MatchingConstructor c]++-- | Matching anything.+matching :: Doc () -> BreadCrumbs+matching d = BreadCrumbs [Matching d]++instance Pretty BreadCrumbs where+  pretty (BreadCrumbs []) = mempty+  pretty (BreadCrumbs bcs) = line <> stack (map pretty bcs)
src/Language/Futhark/TypeChecker/Modules.hs view
@@ -58,19 +58,19 @@ -- | All names defined anywhere in the 'Env'. allNamesInEnv :: Env -> S.Set VName allNamesInEnv (Env vtable ttable stable modtable _names) =-  S.fromList-    ( M.keys vtable-        ++ M.keys ttable-        ++ M.keys stable-        ++ M.keys modtable-    )-    <> mconcat-      ( map allNamesInMTy (M.elems stable)-          ++ map allNamesInMod (M.elems modtable)-          ++ map allNamesInType (M.elems ttable)-      )+  S.unions+    [ M.keysSet vtable,+      M.keysSet ttable,+      M.keysSet stable,+      M.keysSet modtable,+      foldMap allNamesInVal (M.elems vtable),+      foldMap allNamesInMTy (M.elems stable),+      foldMap allNamesInMod (M.elems modtable),+      foldMap allNamesInType (M.elems ttable)+    ]   where     allNamesInType (TypeAbbr _ ps _) = S.fromList $ map typeParamName ps+    allNamesInVal (BoundV ps _) = S.fromList $ map typeParamName ps  allNamesInMod :: Mod -> S.Set VName allNamesInMod (ModEnv env) = allNamesInEnv env@@ -148,8 +148,8 @@             (substituteInMod mod)             (substituteInMTy substs mty) -        substituteInTypeBinding (TypeAbbr l ps (RetType dims t)) =-          TypeAbbr l (map substituteInTypeParam ps) $ RetType dims $ substituteInType t+        substituteInTypeBinding (TypeAbbr l ps t) =+          TypeAbbr l (map substituteInTypeParam ps) $ substituteInRetType t          substituteInTypeParam (TypeParamDim p loc) =           TypeParamDim (substitute p) loc@@ -169,7 +169,8 @@           Arrow als v d1 (substituteInType t1) $ RetType dims $ substituteInType t2          substituteInRetType :: RetTypeBase Size u -> RetTypeBase Size u-        substituteInRetType (RetType ext t) = RetType ext $ substituteInType t+        substituteInRetType (RetType ext t) =+          RetType (map substitute ext) $ substituteInType t          substituteInType :: TypeBase Size u -> TypeBase Size u         substituteInType (Scalar t) = Scalar $ substituteInScalarType t
src/Language/Futhark/TypeChecker/Monad.hs view
@@ -11,6 +11,7 @@     bindSpaced1,     bindIdents,     qualifyTypeVars,+    qualifyTypeVarsWith,     lookupMTy,     lookupImport,     lookupMod,@@ -28,6 +29,7 @@     MonadTypeChecker (..),     TypeState (stateNameSource),     addTySet,+    getTySet,     collectTySet,     usedName,     checkName,@@ -65,6 +67,7 @@ import Control.Monad.Reader import Control.Monad.State.Strict import Data.Either+import Data.IntSet qualified as IS import Data.List (find) import Data.Map.Strict qualified as M import Data.Maybe@@ -176,8 +179,8 @@ data TypeState = TypeState   { stateNameSource :: VNameSource,     stateWarnings :: Warnings,-    -- | Which names have been used.-    stateUsed :: S.Set VName,+    -- | Which names have been used?+    stateUsed :: IS.IntSet,     -- | Known abstract type names.     stateTySet :: TySet,     stateCounter :: Int@@ -255,6 +258,10 @@ addTySet :: TySet -> TypeM () addTySet tys = modify $ \s -> s {stateTySet = tys <> stateTySet s} +-- | Retrieve set of abstract types.+getTySet :: TypeM TySet+getTySet = gets stateTySet+ -- | Run type checking command while accumulating (and returning) all new -- abstract types, then reset to known abstract types afterwards. collectTySet :: TypeM a -> TypeM (a, TySet)@@ -308,10 +315,11 @@   put s {stateCounter = stateCounter s + 1}   pure $ stateCounter s -bindNameMap :: NameMap -> TypeM a -> TypeM a-bindNameMap m = local $ \ctx ->+-- | Run the given action with the name map transformed by the given function.+withNameMap :: (NameMap -> NameMap) -> TypeM a -> TypeM a+withNameMap f = local $ \ctx ->   let env = contextEnv ctx-   in ctx {contextEnv = env {envNameMap = m <> envNameMap env}}+   in ctx {contextEnv = env {envNameMap = f (envNameMap env)}}  -- | Monads that support type checking.  The reason we have this -- internal interface is because we use distinct monads for checking@@ -331,10 +339,10 @@    typeError :: (Located loc) => loc -> Notes -> Doc () -> m a -warnIfUnused :: (Namespace, VName, SrcLoc) -> TypeM ()-warnIfUnused (ns, name, loc) = do+warnIfUnused :: Namespace -> VName -> SrcLoc -> TypeM ()+warnIfUnused ns name loc = do   used <- gets stateUsed-  unless (name `S.member` used || "_" `T.isPrefixOf` nameToText (baseName name)) $+  unless (baseTag name `IS.member` used || "_" `T.isPrefixOf` nameToText (baseName name)) $     warn loc $       "Unused" <+> pretty ns <+> dquotes (prettyName name) <> "." @@ -343,35 +351,32 @@ bindSpaced :: [(Namespace, Name, SrcLoc)] -> ([VName] -> TypeM a) -> TypeM a bindSpaced names body = do   names' <- mapM (\(_, v, _) -> newID v) names-  let mapping = M.fromList $ zip (map (\(ns, v, _) -> (ns, v)) names) $ map qualName names'-  bindNameMap mapping (body names')-    <* mapM_ warnIfUnused [(ns, v, loc) | ((ns, _, loc), v) <- zip names names']+  let ins nm ((ns, v, _), v') = M.insert (ns, v) (qualName v') nm+  withNameMap (\nm -> foldl' ins nm (zip names names')) (body names')+    <* zipWithM_ (\(ns, _, loc) v -> warnIfUnused ns v loc) names names'  -- | Map single source-level name to fresh unique internal names, and -- evaluate a type checker context with the mapping active. bindSpaced1 :: Namespace -> Name -> SrcLoc -> (VName -> TypeM a) -> TypeM a bindSpaced1 ns name loc body = do   name' <- newID name-  let mapping = M.singleton (ns, name) $ qualName name'-  bindNameMap mapping (body name') <* warnIfUnused (ns, name', loc)+  withNameMap (M.insert (ns, name) (qualName name')) (body name')+    <* warnIfUnused ns name' loc  -- | Bind these identifiers in the name map and also check whether -- they have been used. bindIdents :: [IdentBase NoInfo VName t] -> TypeM a -> TypeM a bindIdents idents body = do-  let mapping =-        M.fromList $-          zip-            (map ((Term,) . (baseName . identName)) idents)-            (map (qualName . identName) idents)-  bindNameMap mapping body <* mapM_ warnIfUnused [(Term, v, loc) | Ident v _ loc <- idents]+  let ins nm (Ident v _ _) = M.insert (Term, baseName v) (qualName v) nm+  withNameMap (\nm -> foldl' ins nm idents) body+    <* mapM_ (\(Ident v _ loc) -> warnIfUnused Term v loc) idents  -- | Indicate that this name has been used. This is usually done -- implicitly by other operations, but sometimes we want to make a -- "fake" use to avoid things like top level functions being -- considered unused. usedName :: VName -> TypeM ()-usedName name = modify $ \s -> s {stateUsed = S.insert name $ stateUsed s}+usedName name = modify $ \s -> s {stateUsed = IS.insert (baseTag name) $ stateUsed s}  instance MonadTypeChecker TypeM where   warnings ws =@@ -494,14 +499,35 @@   [VName] ->   TypeBase Size as ->   TypeBase Size as-qualifyTypeVars outer_env orig_except ref_qs = onType (S.fromList orig_except)+qualifyTypeVars = qualifyTypeVarsWith onDim   where+    onDim qual except e = runIdentity $ onDimM except e+      where+        onDimM except' (Var qn typ loc) = pure $ Var (qual except' qn) typ loc+        onDimM except' e' = astMap (identityMapper {mapOnExp = onDimM except'}) e'++-- | Like 'qualifyTypeVars', but generic in the representation of+-- sizes, which are handled by the given function (that is passed the+-- qualification function and the set of names not to qualify).+qualifyTypeVarsWith ::+  forall dim as.+  ((S.Set VName -> QualName VName -> QualName VName) -> S.Set VName -> dim -> dim) ->+  Env ->+  [VName] ->+  [VName] ->+  TypeBase dim as ->+  TypeBase dim as+qualifyTypeVarsWith onDim outer_env orig_except ref_qs+  | null ref_qs = id+  | otherwise = onType (S.fromList orig_except)+  where     onType ::+      forall as'.       S.Set VName ->-      TypeBase Size as ->-      TypeBase Size as+      TypeBase dim as' ->+      TypeBase dim as'     onType except (Array u shape et) =-      Array u (fmap (onDim except) shape) (onScalar except et)+      Array u (fmap (onDim qual except) shape) (onScalar except et)     onType except (Scalar t) =       Scalar $ onScalar except t @@ -520,13 +546,9 @@           Unnamed -> except      onTypeArg except (TypeArgDim d) =-      TypeArgDim $ onDim except d+      TypeArgDim $ onDim qual except d     onTypeArg except (TypeArgType t) =       TypeArgType $ onType except t--    onDim except e = runIdentity $ onDimM except e-    onDimM except (Var qn typ loc) = pure $ Var (qual except qn) typ loc-    onDimM except e = astMap (identityMapper {mapOnExp = onDimM except}) e      qual except (QualName orig_qs name)       | name `elem` except || reachable orig_qs name outer_env =
src/Language/Futhark/TypeChecker/Names.hs view
@@ -22,7 +22,6 @@ import Data.Text qualified as T import Futhark.Util.Pretty import Language.Futhark-import Language.Futhark.Semantic (includeToFilePath) import Language.Futhark.TypeChecker.Monad import Prelude hiding (mod) @@ -119,9 +118,8 @@   v' <- checkValName v loc   case v' of     QualName (q : _) _-      | isIntrinsic q -> do-          me <- askImportName-          unless (isBuiltin (includeToFilePath me)) $+      | isIntrinsic q ->+          unless (isBuiltinLoc loc) $             warn loc "Using intrinsic functions directly can easily crash the compiler or result in wrong code generation."     _ -> pure ()   pure v'@@ -498,10 +496,21 @@   attrs' <- mapM resolveAttrInfo attrs   checkForDuplicateNames tparams params   checkDoNotShadow loc fname-  resolveTypeParams tparams $ \tparams' ->-    resolveParams params $ \params' -> do-      ret' <- traverse resolveTypeExp ret-      body' <- resolveExp body-      bindSpaced1 Term fname loc $ \fname' -> do-        usedName fname'-        pure $ ValBind entry fname' fname_loc ret' NoInfo tparams' params' body' doc attrs' loc+  resolveTypeParams tparams $ \tparams' -> do+    -- Allow self-reference (recursion) only for syntactic functions, i.e.+    -- those with parameters. See Note [Checking recursive functions].+    (fname', params', ret', body') <-+      case params of+        [] -> resolveParams params $ \params' -> do+          ret' <- traverse resolveTypeExp ret+          body' <- resolveExp body+          bindSpaced1 Term fname loc $ \fname' -> do+            usedName fname'+            pure (fname', params', ret', body')+        _ -> bindSpaced1 Term fname loc $ \fname' -> do+          resolveParams params $ \params' -> do+            ret' <- traverse resolveTypeExp ret+            body' <- resolveExp body+            usedName fname'+            pure (fname', params', ret', body')+    pure $ ValBind entry fname' fname_loc ret' NoInfo tparams' params' body' doc attrs' loc
src/Language/Futhark/TypeChecker/Terms.hs view
@@ -6,1694 +6,2248 @@ -- number of built-in language constructs, as well as uniqueness -- types.  This is mostly done in an ad hoc way, and many programs -- will require the programmer to fall back on type annotations.-module Language.Futhark.TypeChecker.Terms-  ( checkOneExp,-    checkSizeExp,-    checkFunDef,-  )-where--import Control.Monad-import Control.Monad.Except-import Control.Monad.Reader-import Control.Monad.State.Strict-import Data.Bifunctor-import Data.Bitraversable-import Data.Char (isAscii)-import Data.Either-import Data.List (delete, find, 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 Data.Text qualified as T-import Futhark.Util (mapAccumLM, nubOrd)-import Futhark.Util.Pretty hiding (space)-import Language.Futhark-import Language.Futhark.Primitive (intByteSize)-import Language.Futhark.Traversals-import Language.Futhark.TypeChecker.Consumption qualified as Consumption-import Language.Futhark.TypeChecker.Match-import Language.Futhark.TypeChecker.Monad hiding (BoundV, lookupMod)-import Language.Futhark.TypeChecker.Terms.Loop-import Language.Futhark.TypeChecker.Terms.Monad-import Language.Futhark.TypeChecker.Terms.Pat-import Language.Futhark.TypeChecker.Types-import Language.Futhark.TypeChecker.Unify-import Prelude hiding (mod)--hasBinding :: Exp -> Bool-hasBinding Lambda {} = True-hasBinding (AppExp LetPat {} _) = True-hasBinding (AppExp LetFun {} _) = True-hasBinding (AppExp Loop {} _) = True-hasBinding (AppExp LetWith {} _) = True-hasBinding (AppExp Match {} _) = True-hasBinding e = isNothing $ astMap m e-  where-    m =-      identityMapper {mapOnExp = \e' -> if hasBinding e' then Nothing else Just e'}--overloadedTypeVars :: Constraints -> Names-overloadedTypeVars = mconcat . map f . M.elems-  where-    f (_, HasFields _ fs _) = mconcat $ map typeVars $ M.elems fs-    f _ = mempty----- Basic checking---- | Determine if the two types are identical, ignoring uniqueness.--- Mismatched dimensions are turned into fresh rigid type variables.--- Causes a 'TypeError' if they fail to match, and otherwise returns--- one of them.-unifyBranchTypes :: SrcLoc -> StructType -> StructType -> TermTypeM (StructType, [VName])-unifyBranchTypes loc t1 t2 =-  onFailure (CheckingBranches t1 t2) $-    unifyMostCommon (mkUsage loc "unification of branch results") t1 t2--unifyBranches :: SrcLoc -> Exp -> Exp -> TermTypeM (StructType, [VName])-unifyBranches loc e1 e2 = do-  e1_t <- expTypeFully e1-  e2_t <- expTypeFully e2-  unifyBranchTypes loc e1_t e2_t--sliceShape ::-  Maybe (SrcLoc, Rigidity) ->-  [DimIndex] ->-  TypeBase Size as ->-  TermTypeM (TypeBase Size as, [VName])-sliceShape r slice t@(Array u (Shape orig_dims) et) =-  runStateT (setDims <$> adjustDims slice orig_dims) []-  where-    setDims [] = stripArray (length orig_dims) t-    setDims dims' = Array u (Shape dims') et--    -- If the result is supposed to be a nonrigid size variable, then-    -- don't bother trying to create non-existential sizes.  This is-    -- necessary to make programs type-check without too much-    -- ceremony; see e.g. tests/inplace5.fut.-    isRigid Rigid {} = True-    isRigid _ = False-    refine_sizes = maybe False (isRigid . snd) r--    sliceSize orig_d i j stride =-      case r of-        Just (loc, Rigid _) -> do-          (d, ext) <--            lift . extSize loc $-              SourceSlice orig_d' (bareExp <$> i) (bareExp <$> j) (bareExp <$> stride)-          modify (maybeToList ext ++)-          pure d-        Just (loc, Nonrigid) ->-          lift $-            flip sizeFromName loc . qualName-              <$> newFlexibleDim (mkUsage loc "size of slice") "slice_dim"-        Nothing -> do-          v <- lift $ newID "slice_anydim"-          modify (v :)-          pure $ sizeFromName (qualName v) mempty-      where-        -- The original size does not matter if the slice is fully specified.-        orig_d'-          | isJust i, isJust j = Nothing-          | otherwise = Just orig_d--    warnIfBinding binds d i j stride size =-      if binds-        then do-          lift . warn (srclocOf size) $-            withIndexLink-              "size-expression-bind"-              "Size expression with binding is replaced by unknown size."-          (:) <$> sliceSize d i j stride-        else pure (size :)--    adjustDims (DimFix {} : idxes') (_ : dims) =-      adjustDims idxes' dims-    -- Pat match some known slices to be non-existential.-    adjustDims (DimSlice i j stride : idxes') (d : dims)-      | refine_sizes,-        maybe True ((== Just 0) . isInt64) i,-        maybe True ((== Just 1) . isInt64) stride = do-          let binds = maybe False hasBinding j-          warnIfBinding binds d i j stride (fromMaybe d j)-            <*> adjustDims idxes' dims-    adjustDims ((DimSlice i j stride) : idxes') (d : dims)-      | refine_sizes,-        Just i' <- i, -- if i ~ 0, previous case-        maybe True ((== Just 1) . isInt64) stride = do-          let j' = fromMaybe d j-              binds = hasBinding j' || hasBinding i'-          warnIfBinding binds d i j stride (sizeMinus j' i')-            <*> adjustDims idxes' dims-    -- stride == -1-    adjustDims ((DimSlice Nothing Nothing stride) : idxes') (d : dims)-      | refine_sizes,-        maybe True ((== Just (-1)) . isInt64) stride =-          (d :) <$> adjustDims idxes' dims-    adjustDims ((DimSlice (Just i) (Just j) stride) : idxes') (d : dims)-      | refine_sizes,-        maybe True ((== Just (-1)) . isInt64) stride = do-          let binds = hasBinding i || hasBinding j-          warnIfBinding binds d (Just i) (Just j) stride (sizeMinus i j)-            <*> adjustDims idxes' dims-    -- existential-    adjustDims ((DimSlice i j stride) : idxes') (d : dims) =-      (:) <$> sliceSize d i j stride <*> adjustDims idxes' dims-    adjustDims _ dims =-      pure dims--    sizeMinus j i =-      AppExp-        ( BinOp-            (qualName (intrinsicVar "-"), mempty)-            sizeBinOpInfo-            (j, Info Nothing)-            (i, Info Nothing)-            mempty-        )-        $ Info-        $ AppRes i64 []-    i64 = Scalar $ Prim $ Signed Int64-    sizeBinOpInfo = Info $ foldFunType [i64, i64] $ RetType [] i64-sliceShape _ _ t = pure (t, [])----- Main checkers--checkAscript ::-  SrcLoc ->-  TypeExp (ExpBase NoInfo VName) VName ->-  ExpBase NoInfo VName ->-  TermTypeM (TypeExp Exp VName, Exp)-checkAscript loc te e = do-  (te', decl_t, _) <- checkTypeExpNonrigid te-  e' <- checkExp e-  e_t <- expTypeFully e'--  onFailure (CheckingAscription (toStruct decl_t) e_t) $-    unify (mkUsage loc "type ascription") (toStruct decl_t) e_t--  pure (te', e')--checkCoerce ::-  SrcLoc ->-  TypeExp (ExpBase NoInfo VName) VName ->-  ExpBase NoInfo VName ->-  TermTypeM (TypeExp Exp VName, StructType, Exp)-checkCoerce loc te e = do-  (te', te_t, ext) <- checkTypeExpNonrigid te-  e' <- checkExp e-  e_t <- expTypeFully e'--  te_t_nonrigid <- makeNonExtFresh ext $ toStruct te_t--  onFailure (CheckingAscription (toStruct te_t) e_t) $-    unify (mkUsage loc "size coercion") e_t te_t_nonrigid--  -- If the type expression had any anonymous dimensions, these will-  -- now be in 'ext'.  Those we keep nonrigid and unify with e_t.-  -- This ensures that 'x :> [1][]i32' does not make the second-  -- dimension unknown.  Use of matchDims is sensible because the-  -- structure of e_t' will be fully known due to the unification, and-  -- te_t because type expressions are complete.-  pure (te', toStruct te_t, e')-  where-    makeNonExtFresh ext = bitraverse onDim pure-      where-        onDim d@(Var v _ _)-          | qualLeaf v `elem` ext = pure d-        onDim d = do-          v <- newTypeName "coerce"-          constrain v . Size Nothing $-            mkUsage-              loc-              "a size coercion where the underlying expression size cannot be determined"-          pure $ sizeFromName (qualName v) (srclocOf d)---- Used to remove unknown sizes from function body types before we--- perform let-generalisation.  This is because if a function is--- inferred to return something of type '[x+y]t' where 'x' or 'y' are--- unknown, we want to turn that into '[z]t', where ''z' is a fresh--- unknown, which is then by let-generalisation turned into--- '?[z].[z]t'.-unscopeUnknown ::-  TypeBase Size u ->-  TermTypeM (TypeBase Size u)-unscopeUnknown t = do-  constraints <- getConstraints-  -- These sizes will be immediately turned into existentials, so we-  -- do not need to care about their location.-  fst <$> sizeFree mempty (expKiller constraints) t-  where-    expKiller _ Var {} = Nothing-    expKiller constraints e =-      S.lookupMin $ S.filter (isUnknown constraints) $ (`S.difference` witnesses) $ fvVars $ freeInExp e-    isUnknown constraints vn-      | Just UnknownSize {} <- snd <$> M.lookup vn constraints = True-    isUnknown _ _ = False-    (witnesses, _) = determineSizeWitnesses $ toStruct t--unscopeType ::-  SrcLoc ->-  [VName] ->-  TypeBase Size as ->-  TermTypeM (TypeBase Size as, [VName])-unscopeType tloc unscoped =-  sizeFree tloc $ find (`elem` unscoped) . fvVars . freeInExp--checkExp :: ExpBase NoInfo VName -> TermTypeM Exp-checkExp (Literal val loc) =-  pure $ Literal val loc-checkExp (Hole _ loc) = do-  t <- newTypeVar loc "t"-  pure $ Hole (Info t) loc-checkExp (StringLit vs loc) =-  pure $ StringLit vs loc-checkExp (IntLit val NoInfo loc) = do-  t <- newTypeVar loc "t"-  mustBeOneOf anyNumberType (mkUsage loc "integer literal") t-  pure $ IntLit val (Info t) loc-checkExp (FloatLit val NoInfo loc) = do-  t <- newTypeVar loc "t"-  mustBeOneOf anyFloatType (mkUsage loc "float literal") t-  pure $ FloatLit val (Info t) loc-checkExp (TupLit es loc) =-  TupLit <$> mapM checkExp es <*> pure loc-checkExp (RecordLit fs loc) =-  RecordLit <$> evalStateT (mapM checkField fs) mempty <*> pure loc-  where-    checkField (RecordFieldExplicit f e rloc) = do-      errIfAlreadySet (unLoc f) rloc-      modify $ M.insert (unLoc f) rloc-      RecordFieldExplicit f <$> lift (checkExp e) <*> pure rloc-    checkField (RecordFieldImplicit name NoInfo rloc) = do-      errIfAlreadySet (baseName (unLoc name)) rloc-      t <- lift $ lookupVar rloc $ qualName $ unLoc name-      modify $ M.insert (baseName (unLoc name)) rloc-      pure $ RecordFieldImplicit name (Info t) rloc--    errIfAlreadySet f rloc = do-      maybe_sloc <- gets $ M.lookup f-      case maybe_sloc of-        Just sloc ->-          lift . typeError rloc mempty $-            "Field"-              <+> dquotes (pretty f)-              <+> "previously defined at"-              <+> pretty (locStrRel rloc sloc)-              <> "."-        Nothing -> pure ()--- No need to type check this, as these are only produced by the--- parser if the elements are monomorphic and all match.-checkExp (ArrayVal vs t loc) =-  pure $ ArrayVal vs t loc-checkExp (ArrayLit all_es _ loc) =-  -- Construct the result type and unify all elements with it.  We-  -- only create a type variable for empty arrays; otherwise we use-  -- the type of the first element.  This significantly cuts down on-  -- the number of type variables generated for pathologically large-  -- multidimensional array literals.-  case all_es of-    [] -> do-      et <- newTypeVar loc "t"-      t <- arrayOfM loc et (Shape [sizeFromInteger 0 mempty])-      pure $ ArrayLit [] (Info t) loc-    e : es -> do-      e' <- checkExp e-      et <- expType e'-      es' <- mapM (unifies "type of first array element" et <=< checkExp) es-      t <- arrayOfM loc et (Shape [sizeFromInteger (genericLength all_es) mempty])-      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-  start_t <- expType start'-  maybe_step' <- case maybe_step of-    Nothing -> pure Nothing-    Just step -> do-      let warning = warn loc "First and second element of range are identical, this will produce an empty array."-      case (start, step) of-        (Literal x _, Literal y _) -> when (x == y) warning-        (Var x_name _ _, Var y_name _ _) -> when (x_name == y_name) warning-        _ -> pure ()-      Just <$> (unifies "use in range expression" start_t =<< checkExp step)--  let unifyRange e = unifies "use in range expression" start_t =<< checkExp e-  end' <- traverse unifyRange end--  end_t <- case end' of-    DownToExclusive e -> expType e-    ToInclusive e -> expType e-    UpToExclusive e -> expType e--  -- Special case some ranges to give them a known size.-  let warnIfBinding binds size =-        if binds-          then do-            warn (srclocOf size) $-              withIndexLink-                "size-expression-bind"-                "Size expression with binding is replaced by unknown size."-            d <- newRigidDim loc RigidRange "range_dim"-            pure (sizeFromName (qualName d) mempty, Just d)-          else pure (size, Nothing)-  (dim, retext) <--    case (isInt64 start', isInt64 <$> maybe_step', end') of-      (Just 0, Just (Just 1), UpToExclusive end'')-        | Scalar (Prim (Signed Int64)) <- end_t ->-            warnIfBinding (hasBinding end'') end''-      (Just 0, Nothing, UpToExclusive end'')-        | Scalar (Prim (Signed Int64)) <- end_t ->-            warnIfBinding (hasBinding end'') end''-      (_, Nothing, UpToExclusive end'')-        | Scalar (Prim (Signed Int64)) <- end_t ->-            warnIfBinding (hasBinding end'' || hasBinding start') $ sizeMinus end'' start'-      (_, Nothing, ToInclusive end'')-        -- No stride means we assume a stride of one.-        | Scalar (Prim (Signed Int64)) <- end_t ->-            warnIfBinding (hasBinding end'' || hasBinding start') $ sizeMinusInc end'' start'-      (Just 1, Just (Just 2), ToInclusive end'')-        | Scalar (Prim (Signed Int64)) <- end_t ->-            warnIfBinding (hasBinding end'') end''-      _ -> do-        d <- newRigidDim loc RigidRange "range_dim"-        pure (sizeFromName (qualName d) mempty, Just d)--  t <- arrayOfM loc start_t (Shape [dim])-  let res = AppRes t (maybeToList retext)--  pure $ AppExp (Range start' maybe_step' end' loc) (Info res)-  where-    i64 = Scalar $ Prim $ Signed Int64-    mkBinOp op t x y =-      AppExp-        ( BinOp-            (qualName (intrinsicVar op), mempty)-            sizeBinOpInfo-            (x, Info Nothing)-            (y, Info Nothing)-            mempty-        )-        (Info $ AppRes t [])-    mkSub = mkBinOp "-" i64-    mkAdd = mkBinOp "+" i64-    sizeMinus j i = j `mkSub` i-    sizeMinusInc j i = (j `mkSub` i) `mkAdd` sizeFromInteger 1 mempty-    sizeBinOpInfo = Info $ foldFunType [i64, i64] $ RetType [] i64-checkExp (Ascript e te loc) = do-  (te', e') <- checkAscript loc te e-  pure $ Ascript e' te' loc-checkExp (Coerce e te NoInfo loc) = do-  (te', te_t, e') <- checkCoerce loc te e-  t <- expTypeFully e'-  t' <- matchDims (const . const pure) t te_t-  pure $ Coerce e' te' (Info t') loc-checkExp (AppExp (BinOp (op, oploc) NoInfo (e1, _) (e2, _) loc) NoInfo) = do-  ftype <- lookupVar oploc op-  e1' <- checkExp e1-  e2' <- checkExp e2--  -- Note that the application to the first operand cannot fix any-  -- existential sizes, because it must by necessity be a function.-  (_, rt, p1_ext, _) <- checkApply loc (Just op, 0) ftype e1'-  (_, rt', p2_ext, retext) <- checkApply loc (Just op, 1) rt e2'--  pure $-    AppExp-      ( BinOp-          (op, oploc)-          (Info ftype)-          (e1', Info p1_ext)-          (e2', Info p2_ext)-          loc-      )-      (Info (AppRes rt' retext))-checkExp (Project k e NoInfo loc) = do-  e' <- checkExp e-  t <- expType e'-  kt <- mustHaveField (mkUsage loc $ docText $ "projection of field " <> dquotes (pretty k)) k t-  pure $ Project k e' (Info kt) loc-checkExp (AppExp (If e1 e2 e3 loc) _) = do-  e1' <- checkExp e1-  e2' <- checkExp e2-  e3' <- checkExp e3--  let bool = Scalar $ Prim Bool-  e1_t <- expType e1'-  onFailure (CheckingRequired [bool] e1_t) $-    unify (mkUsage e1' "use as 'if' condition") bool e1_t--  (brancht, retext) <- unifyBranches loc e2' e3'--  zeroOrderType-    (mkUsage loc "returning value of this type from 'if' expression")-    "type returned from branch"-    brancht--  pure $ AppExp (If e1' e2' e3' loc) (Info $ AppRes brancht retext)-checkExp (Parens e loc) =-  Parens <$> checkExp e <*> pure loc-checkExp (QualParens (modname, modnameloc) e loc) = do-  mod <- lookupMod modname-  case mod of-    ModEnv env -> local (`withEnv` env) $ do-      e' <- checkExp e-      pure $ QualParens (modname, modnameloc) e' loc-    ModFun {} ->-      typeError loc mempty . withIndexLink "module-is-parametric" $-        "Module" <+> pretty modname <+> " is a parametric module."-checkExp (Var qn NoInfo loc) = do-  t <- lookupVar loc qn-  pure $ Var qn (Info t) loc-checkExp (Negate arg loc) = do-  arg' <- require "numeric negation" anyNumberType =<< checkExp arg-  pure $ Negate arg' loc-checkExp (Not arg loc) = do-  arg' <- require "logical negation" (Bool : anyIntType) =<< checkExp arg-  pure $ Not arg' loc-checkExp (AppExp (Apply fe args loc) NoInfo) = do-  fe' <- checkExp fe-  args' <- mapM (checkExp . 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-      (_, rt, argext, exts) <- checkApply loc (fname, i) t arg'-      pure-        ( (i + 1, all_exts <> exts, rt),-          (Info argext, arg')-        )-checkExp (AppExp (LetPat sizes pat e body loc) _) = do-  e' <- checkExp e--  -- Not technically an ascription, but we want the pattern to have-  -- exactly the type of 'e'.-  t <- expType e'-  bindingSizes sizes . incLevel . bindingPat sizes pat t $ \pat' -> do-    body' <- incLevel $ checkExp body-    body_t <- expTypeFully body'--    -- If the bound expression is of type i64, then we replace the-    -- pattern name with the expression in the type of the body.-    -- Otherwise, we need to come up with unknown sizes for the-    -- sizes going out of scope.-    t' <- normType t -- Might be overloaded integer until now.-    (body_t', retext) <--      case (t', patNames pat') of-        (Scalar (Prim (Signed Int64)), [v])-          | not $ hasBinding e' -> do-              let f x = if x == v then Just (ExpSubst e') else Nothing-              pure (applySubst f body_t, [])-        _ ->-          unscopeType loc (map sizeName sizes <> patNames pat') body_t--    pure $-      AppExp-        (LetPat sizes (fmap toStruct pat') e' body' loc)-        (Info $ AppRes body_t' retext)-checkExp (AppExp (LetFun name (tparams, params, maybe_retdecl, NoInfo, e) body loc) _) = do-  (tparams', params', maybe_retdecl', rettype, e') <--    checkBinding (fst name, maybe_retdecl, tparams, params, e, loc)--  let entry = BoundV tparams' $ funType params' rettype-      bindF scope =-        scope-          { scopeVtable = M.insert (fst name) entry $ scopeVtable scope-          }-  body' <- localScope bindF $ checkExp body--  (body_t, ext) <- unscopeType loc [fst name] =<< expTypeFully body'--  pure $-    AppExp-      ( LetFun-          name-          (tparams', params', maybe_retdecl', Info rettype, e')-          body'-          loc-      )-      (Info $ AppRes body_t ext)-checkExp (AppExp (LetWith dest src steps ve body loc) _) = do-  src' <- checkIdent src-  src_t <- normTypeFully $ unInfo $ identType src'--  let onlyFields = all isField steps--  if onlyFields-    then do-      ve' <- checkExp ve-      ve_t <- expType ve'-      updated_t <- updateFieldPath src (fieldNames steps) ve_t src_t-      steps' <- mapM checkFieldStep steps--      bindingIdent dest updated_t $ \dest' -> do-        body' <- checkExp body-        (body_t, ext) <- unscopeType loc [identName dest'] =<< expTypeFully body'-        pure $ AppExp (LetWith dest' src' steps' ve' body' loc) (Info $ AppRes body_t ext)-    else do-      (steps', target_t) <- checkUpdateSteps loc src_t steps-      ve' <- unifies "type of update target" target_t =<< checkExp ve--      src_t' <- normTypeFully $ unInfo $ identType src'-      bindingIdent dest src_t' $ \dest' -> do-        body' <- checkExp body-        (body_t, ext) <- unscopeType loc [identName dest'] =<< expTypeFully body'-        pure $ AppExp (LetWith dest' src' steps' ve' body' loc) (Info $ AppRes body_t ext)-  where-    isField UpdateStepField {} = True-    isField _ = False--    fieldNames = map (\(UpdateStepField f) -> f)--    checkFieldStep (UpdateStepField f) = pure $ UpdateStepField f-    checkFieldStep _ = error "impossible"---- Record updates are a bit hacky, because we do not have row typing--- (yet?).  For now, we only permit record updates where we know the--- full type up to the field we are updating.-checkExp (Update src steps ve NoInfo loc) = do-  src' <- checkExp src-  src_t <- expTypeFully src'-  let onlyFields = all isField steps-  if onlyFields-    then do-      ve' <- checkExp ve-      ve_t <- expType ve'-      updated_t <- updateFieldPath src (fieldNames steps) ve_t src_t-      steps' <- mapM checkFieldStep steps-      pure $ Update src' steps' ve' (Info updated_t) loc-    else do-      (steps', target_t) <- checkUpdateSteps loc src_t steps-      ve' <- unifies "type of update target" target_t =<< checkExp ve-      src_t' <- expTypeFully src'-      pure $ Update src' steps' ve' (Info src_t') loc-  where-    isField UpdateStepField {} = True-    isField _ = False--    fieldNames = map (\(UpdateStepField f) -> f)--    checkFieldStep (UpdateStepField f) = pure $ UpdateStepField f-    checkFieldStep _ = error "impossible"-checkExp (AppExp (Index e slice loc) _) = do-  slice' <- checkSlice slice-  (t, _) <- newArrayType (mkUsage' loc) "e" $ sliceDims slice'-  e' <- unifies "being indexed at" t =<< checkExp e-  -- XXX, the RigidSlice here will be overridden in sliceShape with a proper value.-  (t', retext) <--    sliceShape (Just (loc, Rigid (RigidSlice Nothing ""))) slice'-      =<< expTypeFully e'--  pure $ AppExp (Index e' slice' loc) (Info $ AppRes t' retext)-checkExp (Assert e1 e2 NoInfo loc) = do-  e1' <- require "being asserted" [Bool] =<< checkExp e1-  e2' <- checkExp e2-  pure $ Assert e1' e2' (Info (prettyText e1)) loc-checkExp (Lambda params body rettype_te NoInfo loc) = do-  (params', body', rettype', RetType dims ty) <--    incLevel . bindingParams [] params $ \params' -> do-      rettype_checked <- traverse checkTypeExpNonrigid rettype_te-      let declared_rettype =-            case rettype_checked of-              Just (_, st, _) -> Just st-              Nothing -> Nothing-      body' <- checkFunBody params' body declared_rettype loc-      body_t <- expTypeFully body'--      params'' <- mapM updateTypes params'--      (rettype', rettype_st) <--        case rettype_checked of-          Just (te, st, ext) ->-            pure (Just te, RetType ext st)-          Nothing -> do-            ret <- inferReturnSizes params'' $ toRes Nonunique body_t-            pure (Nothing, ret)--      pure (params'', body', rettype', rettype_st)--  verifyFunctionParams Nothing params'--  (ty', dims') <- unscopeType loc dims ty--  pure $ Lambda params' body' rettype' (Info (RetType dims' ty')) loc-  where-    -- Inferring the sizes of the return type of a lambda is a lot-    -- like let-generalisation.  We wish to remove any rigid sizes-    -- that were created when checking the body, except for those that-    -- are visible in types that existed before we entered the body,-    -- are parameters, or are used in parameters.-    inferReturnSizes params' ret = do-      cur_lvl <- curLevel-      let named (Named x, _, _) = Just x-          named (Unnamed, _, _) = Nothing-          param_names = mapMaybe (named . patternParam) params'-          pos_sizes =-            sizeNamesPos $ funType params' $ RetType [] ret-          hide k (lvl, _) =-            lvl >= cur_lvl && k `notElem` param_names && k `S.notMember` pos_sizes--      hidden_sizes <--        S.fromList . M.keys . M.filterWithKey hide <$> getConstraints--      let onDim name-            | name `S.member` hidden_sizes = S.singleton name-          onDim _ = mempty--      pure $ RetType (S.toList $ foldMap onDim $ fvVars $ freeInType ret) ret-checkExp (OpSection op _ loc) = do-  ftype <- lookupVar loc op-  pure $ OpSection op (Info ftype) loc-checkExp (OpSectionLeft op _ e _ _ loc) = do-  ftype <- lookupVar loc op-  e' <- checkExp e-  (t1, rt, argext, retext) <- checkApply loc (Just op, 0) ftype e'-  case (ftype, rt) of-    (Scalar (Arrow _ m1 d1 _ _), Scalar (Arrow _ m2 d2 t2 rettype)) ->-      pure $-        OpSectionLeft-          op-          (Info ftype)-          e'-          (Info (m1, toParam d1 t1, argext), Info (m2, toParam d2 t2))-          (Info rettype, Info retext)-          loc-    _ ->-      typeError loc mempty $-        "Operator section with invalid operator of type" <+> pretty ftype-checkExp (OpSectionRight op _ e _ NoInfo loc) = do-  ftype <- lookupVar loc op-  e' <- checkExp e-  case ftype of-    Scalar (Arrow _ m1 d1 t1 (RetType [] (Scalar (Arrow _ m2 d2 t2 (RetType dims2 ret))))) -> do-      (t2', arrow', argext, _) <--        checkApply-          loc-          (Just op, 1)-          (Scalar $ Arrow mempty m2 d2 t2 $ RetType [] $ Scalar $ Arrow Nonunique m1 d1 t1 $ RetType dims2 ret)-          e'-      case arrow' of-        Scalar (Arrow _ _ _ t1' (RetType dims2' ret')) ->-          pure $-            OpSectionRight-              op-              (Info ftype)-              e'-              (Info (m1, toParam d1 t1'), Info (m2, toParam d2 t2', argext))-              (Info $ RetType dims2' ret')-              loc-        _ -> error $ "OpSectionRight: impossible type\n" <> prettyString arrow'-    _ ->-      typeError loc mempty $-        "Operator section with invalid operator of type" <+> pretty ftype-checkExp (UpdateSection steps NoInfo loc) = do-  a <- newTypeVar loc "a"-  (steps', b, retext) <- checkSectionSteps a steps-  let ft = Scalar $ Arrow mempty Unnamed Observe a $ RetType retext $ toRes Nonunique b-  pure $ UpdateSection steps' (Info ft) loc-  where-    checkSectionSteps t [] =-      pure ([], t, [])-    checkSectionSteps t (step : rest) =-      case step of-        UpdateStepField f -> do-          t' <- mustHaveField (mkUsage loc "projection at") f t-          (rest', target_t, retext) <- checkSectionSteps t' rest-          pure (UpdateStepField f : rest', target_t, retext)-        UpdateStepSlice slice -> do-          slice' <- checkSlice slice-          (arr_t, _) <- newArrayType (mkUsage' loc) "e" $ sliceDims slice'-          unify (mkUsage loc "type of section indexing") arr_t t-          (t', retext) <- sliceShape Nothing slice' =<< normTypeFully arr_t-          (rest', target_t, retext_rest) <- checkSectionSteps t' rest-          pure (UpdateStepSlice slice' : rest', target_t, retext <> retext_rest)-checkExp (AppExp (Loop _ mergepat loopinit form loopbody loc) _) = do-  ((sparams, mergepat', loopinit', form', loopbody'), appres) <--    checkLoop checkExp (mergepat, loopinit, form, loopbody) loc-  pure $-    AppExp-      (Loop sparams mergepat' loopinit' form' loopbody' loc)-      (Info appres)-checkExp (Constr name es NoInfo loc) = do-  t <- newTypeVar loc "t"-  es' <- mapM checkExp es-  ets <- mapM expType es'-  mustHaveConstr (mkUsage loc "use of constructor") name t ets-  pure $ Constr name es' (Info t) loc-checkExp (AppExp (Match e cs loc) _) = do-  e' <- checkExp e-  mt <- expType e'-  (cs', t, retext) <- checkCases mt cs-  zeroOrderType-    (mkUsage loc "being returned 'match'")-    "type returned from pattern match"-    t-  pure $ AppExp (Match e' cs' loc) (Info $ AppRes t retext)-checkExp (Attr info e loc) =-  Attr <$> checkAttr info <*> checkExp e <*> pure loc--updateFieldPath ::-  (Pretty a, Located a) =>-  a ->-  [Name] ->-  StructType ->-  StructType ->-  TermTypeM StructType-updateFieldPath src all_fs ve_t = recurse [] all_fs-  where-    recurse seen [] t = do-      (t', _) <- allDimsFreshInType usage Nonrigid "any" t-      onFailure (CheckingRecordUpdate seen t' ve_t) $-        unify usage t' ve_t-      pure ve_t-      where-        usage = mkUsage (locOf src) "record update"-    recurse seen (f : fs) (Scalar (Record m))-      | Just f_t <- M.lookup f m = do-          f_t' <- recurse (seen ++ [f]) fs f_t-          pure $ Scalar $ Record $ M.insert f f_t' m-    recurse _ _ _ =-      typeError (locOf src) mempty . withIndexLink "record-type-not-known" $-        "Full type of"-          </> indent 2 (pretty src)-          </> textwrap " is not known at this point.  Add a type annotation to the original record to disambiguate."--checkUpdateSteps ::-  SrcLoc ->-  StructType ->-  [UpdateStep NoInfo VName] ->-  TermTypeM ([UpdateStep Info VName], StructType)-checkUpdateSteps _ t [] =-  pure ([], t)-checkUpdateSteps loc t (step : rest) =-  case step of-    UpdateStepSlice slice -> do-      slice' <- checkSlice slice-      (arr_t, _) <- newArrayType (mkUsage' loc) "update_path_src" $ sliceDims slice'-      unify (mkUsage loc "type of update path indexing") arr_t t-      (elem_t, _) <- sliceShape (Just (loc, Nonrigid)) slice' =<< normTypeFully arr_t-      (rest', target_t) <- checkUpdateSteps loc elem_t rest-      pure (UpdateStepSlice slice' : rest', target_t)-    UpdateStepField f -> do-      t' <- normTypeFully t-      f_t <- mustHaveField (mkUsage loc "record update path") f t'-      (rest', target_t) <- checkUpdateSteps loc f_t rest-      pure (UpdateStepField f : rest', target_t)--checkCases ::-  StructType ->-  NE.NonEmpty (CaseBase NoInfo VName) ->-  TermTypeM (NE.NonEmpty (CaseBase Info VName), StructType, [VName])-checkCases mt rest_cs =-  case NE.uncons rest_cs of-    (c, Nothing) -> do-      (c', t, retext) <- checkCase mt c-      pure (NE.singleton c', t, retext)-    (c, Just cs) -> do-      ((c', c_t, _), (cs', cs_t, _)) <--        (,) <$> checkCase mt c <*> checkCases mt cs-      (brancht, retext) <- unifyBranchTypes (srclocOf c) c_t cs_t-      pure (NE.cons c' cs', brancht, retext)--checkCase ::-  StructType ->-  CaseBase NoInfo VName ->-  TermTypeM (CaseBase Info VName, StructType, [VName])-checkCase mt (CasePat p e loc) =-  bindingPat [] p mt $ \p' -> do-    e' <- checkExp e-    e_t <- expTypeFully e'-    (e_t', retext) <- unscopeType loc (patNames p') e_t-    pure (CasePat (fmap toStruct p') e' loc, e_t', retext)---- | An unmatched pattern. Used in in the generation of--- unmatched pattern warnings by the type checker.-data Unmatched p-  = UnmatchedNum p [PatLit]-  | UnmatchedBool p-  | UnmatchedConstr p-  | Unmatched p-  deriving (Functor, Show)--instance Pretty (Unmatched (Pat StructType)) where-  pretty um = case um of-    (UnmatchedNum p nums) -> pretty' p <+> "where p is not one of" <+> pretty nums-    (UnmatchedBool p) -> pretty' p-    (UnmatchedConstr p) -> pretty' p-    (Unmatched p) -> pretty' p-    where-      pretty' (PatAscription p t _) = pretty p <> ":" <+> pretty t-      pretty' (PatParens p _) = parens $ pretty' p-      pretty' (PatAttr _ p _) = parens $ pretty' p-      pretty' (Id v _ _) = prettyName v-      pretty' (TuplePat pats _) = parens $ commasep $ map pretty' pats-      pretty' (RecordPat fs _) = braces $ commasep $ map ppField fs-        where-          ppField (L _ name, t) = pretty (nameToString name) <> equals <> pretty' t-      pretty' Wildcard {} = "_"-      pretty' (PatLit e _ _) = pretty e-      pretty' (PatConstr n _ ps _) = "#" <> pretty n <+> sep (map pretty' ps)--checkIdent :: IdentBase NoInfo VName StructType -> TermTypeM (Ident StructType)-checkIdent (Ident name _ loc) = do-  vt <- lookupVar loc $ qualName name-  pure $ Ident name (Info vt) loc--checkSlice :: SliceBase NoInfo VName -> TermTypeM [DimIndex]-checkSlice = mapM checkDimIndex-  where-    checkDimIndex (DimFix i) = do-      DimFix <$> (require "use as index" anySignedType =<< checkExp i)-    checkDimIndex (DimSlice i j s) =-      DimSlice <$> check i <*> check j <*> check s--    check =-      maybe (pure Nothing) $-        fmap Just . unifies "use as index" (Scalar $ Prim $ Signed Int64) <=< checkExp---- The number of dimensions affected by this slice (so the minimum--- rank of the array we are slicing).-sliceDims :: [DimIndex] -> Int-sliceDims = length--instantiateDimsInReturnType ::-  SrcLoc ->-  Maybe (QualName VName) ->-  ResRetType ->-  TermTypeM (ResType, [VName])-instantiateDimsInReturnType loc fname (RetType dims t)-  | null dims =-      pure (t, mempty)-  | otherwise = do-      dims' <- mapM new dims-      pure (first (onDim $ zip dims $ map (ExpSubst . (`sizeFromName` loc) . qualName) dims') t, dims')-  where-    new =-      newRigidDim loc (RigidRet fname)-        . nameFromText-        . T.takeWhile isAscii-        . baseText-    onDim dims' = applySubst (`lookup` dims')---- Some information about the function/operator we are trying to--- apply, and how many arguments it has previously accepted.  Used for--- generating nicer type errors.-type ApplyOp = (Maybe (QualName VName), Int)---- | Extract all those names that are bound inside the type.-boundInsideType :: TypeBase Size as -> S.Set VName-boundInsideType (Array _ _ t) = boundInsideType (Scalar t)-boundInsideType (Scalar Prim {}) = mempty-boundInsideType (Scalar (TypeVar _ _ targs)) = foldMap f targs-  where-    f (TypeArgType t) = boundInsideType t-    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))) =-  pn' <> boundInsideType t1 <> S.fromList dims <> boundInsideType t2-  where-    pn' = case pn of-      Unnamed -> mempty-      Named v -> S.singleton v---- Returns the sizes of the immediate type produced,--- the sizes of parameter types, and the sizes of return types.-dimUses :: TypeBase Size u -> (Names, Names)-dimUses = flip execState mempty . traverseDims f-  where-    f bound pos e =-      case pos of-        PosImmediate ->-          modify ((fvVars fv, mempty) <>)-        PosParam ->-          modify ((mempty, fvVars fv) <>)-        PosReturn -> pure ()-      where-        fv = freeInExp e `freeWithout` bound--checkApply ::-  SrcLoc ->-  ApplyOp ->-  StructType ->-  Exp ->-  TermTypeM (StructType, StructType, Maybe VName, [VName])-checkApply loc (fname, _) (Scalar (Arrow _ pname _ tp1 tp2)) argexp = do-  let argtype = typeOf argexp-  onFailure (CheckingApply fname argexp tp1 argtype) $ do-    unify (mkUsage argexp "use as function argument") tp1 argtype--    -- Perform substitutions of instantiated variables in the types.-    (tp2', ext) <- instantiateDimsInReturnType loc fname =<< normTypeFully tp2-    argtype' <- normTypeFully argtype--    -- Check whether this would produce an impossible return type.-    let (tp2_produced_dims, tp2_paramdims) = dimUses tp2'-        problematic = S.fromList ext <> boundInsideType argtype'-        problem = any (`S.member` problematic) (tp2_paramdims `S.difference` tp2_produced_dims)-    when (not (S.null problematic) && problem) $ do-      typeError loc mempty . withIndexLink "existential-param-ret" $-        "Existential size would appear in function parameter of return type:"-          </> indent 2 (pretty (RetType ext tp2'))-          </> textwrap "This is usually because a higher-order function is used with functional arguments that return existential sizes or locally named sizes, which are then used as parameters of other function arguments."--    (argext, tp2'') <--      case pname of-        Named pname'-          | S.member pname' (fvVars $ freeInType tp2') ->-              if hasBinding argexp-                then do-                  warn (srclocOf argexp) $-                    withIndexLink-                      "size-expression-bind"-                      "Size expression with binding is replaced by unknown size."-                  d <- newRigidDim argexp (RigidArg fname $ prettyTextOneLine $ bareExp argexp) "n"-                  let parsubst v =-                        if v == pname'-                          then Just $ ExpSubst $ sizeFromName (qualName d) $ srclocOf argexp-                          else Nothing-                  pure (Just d, applySubst parsubst $ toStruct tp2')-                else-                  let parsubst v =-                        if v == pname'-                          then Just $ ExpSubst $ fromMaybe argexp $ stripExp argexp-                          else Nothing-                   in pure (Nothing, applySubst parsubst $ toStruct tp2')-        _ -> pure (Nothing, toStruct tp2')--    pure (tp1, tp2'', argext, ext)-checkApply loc fname tfun@(Scalar TypeVar {}) arg = do-  tv <- newTypeVar loc "b"-  unify (mkUsage loc "use as function") tfun $-    Scalar (Arrow mempty Unnamed Observe (typeOf arg) $ RetType [] $ paramToRes tv)-  tfun' <- normType tfun-  checkApply loc fname tfun' arg-checkApply loc (fname, prev_applied) ftype argexp = do-  let fname' = maybe "expression" (dquotes . pretty) fname--  typeError loc mempty $-    if prev_applied == 0-      then-        "Cannot apply"-          <+> fname'-          <+> "as function, as it has type:"-          </> indent 2 (pretty ftype)-      else-        "Cannot apply"-          <+> fname'-          <+> "to argument #"-          <> pretty (prev_applied + 1)-            <+> dquotes (shorten $ group $ pretty argexp)-          <> ","-            </> "as"-            <+> fname'-            <+> "only takes"-            <+> pretty prev_applied-            <+> arguments-          <> "."-  where-    arguments-      | prev_applied == 1 = "argument"-      | otherwise = "arguments"---- | Type-check a single expression in isolation.  This expression may--- turn out to be polymorphic, in which case the list of type--- parameters will be non-empty.-checkOneExp :: ExpBase NoInfo VName -> TypeM ([TypeParam], Exp)-checkOneExp e = runTermTypeM checkExp $ do-  e' <- checkExp e-  (tparams, _, RetType _ t') <--    letGeneralise (nameFromString "<exp>") (srclocOf e) [] [] $-      toRes Nonunique $-        typeOf e'-  fixOverloadedTypes $ typeVars t'-  e'' <- normTypeFully e'-  localChecks e''-  causalityCheck e''-  pure (tparams, e'')---- | Type-check a single size expression in isolation.  This expression may--- turn out to be polymorphic, in which case it is unified with i64.-checkSizeExp :: ExpBase NoInfo VName -> TypeM Exp-checkSizeExp e = runTermTypeM checkExp $ do-  e' <- checkExp e-  let t = typeOf e'-  when (hasBinding e') $-    typeError (srclocOf e') mempty . withIndexLink "size-expression-bind" $-      "Size expression with binding is forbidden."-  unify (mkUsage e' "Size expression") t (Scalar (Prim (Signed Int64)))-  normTypeFully e'---- Verify that all sum type constructors and empty array literals have--- a size that is known (rigid or a type parameter).  This is to--- ensure that we can actually determine their shape at run-time.-causalityCheck :: Exp -> TermTypeM ()-causalityCheck binding_body = do-  constraints <- getConstraints--  let checkCausality what known t loc-        | (d, dloc) : _ <--            mapMaybe (unknown constraints known) $-              S.toList (fvVars $ freeInType t) =-            Just $ lift $ causality what (locOf loc) d dloc t-        | otherwise = Nothing--      checkParamCausality known p =-        checkCausality (pretty p) known (patternType p) (locOf p)--      collectingNewKnown = lift . flip execStateT mempty--      onExp ::-        S.Set VName ->-        Exp ->-        StateT (S.Set VName) (Either TypeError) Exp--      onExp known (Var v (Info t) loc)-        | Just bad <- checkCausality (dquotes (pretty v)) known t loc =-            bad-      onExp known (UpdateSection _ (Info t) loc)-        | Just bad <- checkCausality "projection section" known t loc =-            bad-      onExp known (OpSectionRight _ (Info t) _ _ _ loc)-        | Just bad <- checkCausality "operator section" known t loc =-            bad-      onExp known (OpSectionLeft _ (Info t) _ _ _ loc)-        | Just bad <- checkCausality "operator section" known t loc =-            bad-      onExp known (ArrayLit [] (Info t) loc)-        | Just bad <- checkCausality "empty array" known t loc =-            bad-      onExp known (Hole (Info t) loc)-        | Just bad <- checkCausality "hole" known t loc =-            bad-      onExp known e@(Lambda params body _ _ _)-        | bad : _ <- mapMaybe (checkParamCausality known) params =-            bad-        | otherwise = do-            -- Existentials coming into existence in the lambda body-            -- are not known outside of it.-            void $ collectingNewKnown $ onExp known body-            pure e-      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 (Match scrutinee cs _) (Info res)) = do-        new_known <- collectingNewKnown $ onExp known scrutinee-        void $ recurse (new_known <> known) cs-        modify ((new_known <> S.fromList (appResExt res)) <>)-        pure e-      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 <- collectingNewKnown $ onExp known' x-            void $ seqArgs (new_known <> known') xs-            modify ((new_known <> S.fromList (maybeToList p)) <>)-      onExp known e@(Constr v args (Info t) loc) = do-        seqArgs known args-        pure e-        where-          seqArgs known' []-            | Just bad <- checkCausality (dquotes ("#" <> pretty v)) known' t loc =-                bad-            | otherwise =-                pure ()-          seqArgs known' (x : xs) = do-            new_known <- collectingNewKnown $ onExp known' x-            void $ seqArgs (new_known <> known') xs-            modify (new_known <>)-      onExp-        known-        e@(AppExp (BinOp (f, floc) ft (x, Info xp) (y, Info yp) _) (Info res)) = do-          args_known <--            collectingNewKnown $ sequencePoint known x y $ catMaybes [xp, yp]-          void $ onExp (args_known <> known) (Var f ft floc)-          modify ((args_known <> S.fromList (appResExt res)) <>)-          pure e-      onExp known e@(AppExp e' (Info res)) = do-        recurse known e'-        modify (<> S.fromList (appResExt res))-        pure e-      onExp known e = do-        recurse known e-        pure e--      recurse known = void . astMap mapper-        where-          mapper = identityMapper {mapOnExp = onExp known}--      sequencePoint known x y ext = do-        new_known <- collectingNewKnown $ onExp known x-        void $ onExp (new_known <> known) y-        modify ((new_known <> S.fromList ext) <>)--  either throwError (const $ pure ()) $-    evalStateT (onExp mempty binding_body) mempty-  where-    unknown constraints known v = do-      guard $ v `S.notMember` known-      loc <- case snd <$> M.lookup v constraints of-        Just (UnknownSize loc _) -> Just loc-        _ -> Nothing-      pure (v, loc)--    causality what loc d dloc t =-      Left . TypeError loc mempty . withIndexLink "causality-check" $-        "Causality check: size"-          <+> dquotes (prettyName d)-          <+> "needed for type of"-          <+> what-          <> colon-            </> indent 2 (pretty t)-            </> "But"-            <+> dquotes (prettyName d)-            <+> "is computed at"-            <+> pretty (locStrRel loc dloc)-          <> "."-            </> ""-            </> "Hint:"-            <+> align-              ( textwrap "Bind the expression producing"-                  <+> dquotes (prettyName d)-                  <+> "with 'let' beforehand."-              )--mustBeIrrefutable :: (MonadTypeChecker f) => Pat StructType -> f ()-mustBeIrrefutable p = do-  case unmatched [p] of-    [] -> pure ()-    ps' ->-      typeError p mempty . withIndexLink "refutable-pattern" $-        "Refutable pattern not allowed here.\nUnmatched cases:"-          </> indent 2 (stack (map pretty ps'))---- | Traverse the expression, emitting warnings and errors for various--- problems:------ * Unmatched cases.------ * If any of the literals overflow their inferred types. Note:---  currently unable to detect float underflow (such as 1e-400 -> 0)-localChecks :: Exp -> TermTypeM ()-localChecks = void . check-  where-    check e@(AppExp (Match _ cs loc) _) = do-      let ps = fmap (\(CasePat p _ _) -> p) cs-      case unmatched $ NE.toList ps of-        [] -> recurse e-        ps' ->-          typeError loc mempty . withIndexLink "unmatched-cases" $-            "Unmatched cases in match expression:"-              </> indent 2 (stack (map pretty ps'))-    check e@(AppExp (LetPat _ p _ _ _) _) =-      mustBeIrrefutable p *> recurse e-    check e@(Lambda ps _ _ _ _) =-      mapM_ (mustBeIrrefutable . fmap toStruct) ps *> recurse e-    check e@(AppExp (LetFun _ (_, ps, _, _, _) _ _) _) =-      mapM_ (mustBeIrrefutable . fmap toStruct) ps *> recurse e-    check e@(AppExp (Loop _ p _ form _ _) _) = do-      mustBeIrrefutable (fmap toStruct p)-      case form of-        ForIn form_p _ -> mustBeIrrefutable form_p-        _ -> pure ()-      recurse e-    check e@(IntLit x ty loc) =-      e <$ case ty of-        Info (Scalar (Prim t)) -> errorBounds (inBoundsI x t) x t loc-        _ -> error "Inferred type of int literal is not a number"-    check e@(FloatLit x ty loc) =-      e <$ case ty of-        Info (Scalar (Prim (FloatType t))) -> errorBounds (inBoundsF x t) x t loc-        _ -> error "Inferred type of float literal is not a float"-    check e@(Negate (IntLit x ty loc1) loc2) =-      e <$ case ty of-        Info (Scalar (Prim t)) -> errorBounds (inBoundsI (-x) t) (-x) t (loc1 <> loc2)-        _ -> error "Inferred type of int literal is not a number"-    check e@(AppExp (BinOp (QualName [] v, _) _ (x, _) _ loc) _)-      | baseName v == "==",-        Array {} <- typeOf x,-        isIntrinsic v = do-          warn loc $-            textwrap-              "Comparing arrays with \"==\" is deprecated and will stop working in a future revision of the language."-          recurse e-    check e = recurse e-    recurse = astMap identityMapper {mapOnExp = check}--    bitWidth ty = 8 * intByteSize ty :: Int--    inBoundsI x (Signed t) = x >= -2 ^ (bitWidth t - 1) && x < 2 ^ (bitWidth t - 1)-    inBoundsI x (Unsigned t) = x >= 0 && x < 2 ^ bitWidth t-    inBoundsI x (FloatType Float16) = not $ isInfinite (fromIntegral x :: Half)-    inBoundsI x (FloatType Float32) = not $ isInfinite (fromIntegral x :: Float)-    inBoundsI x (FloatType Float64) = not $ isInfinite (fromIntegral x :: Double)-    inBoundsI _ Bool = error "Inferred type of int literal is not a number"-    inBoundsF x Float16 = not $ isInfinite (realToFrac x :: Float)-    inBoundsF x Float32 = not $ isInfinite (realToFrac x :: Float)-    inBoundsF x Float64 = not $ isInfinite x--    errorBounds inBounds x ty loc =-      unless inBounds $-        typeError loc mempty . withIndexLink "literal-out-of-bounds" $-          "Literal "-            <> pretty x-            <> " out of bounds for inferred type "-            <> pretty ty-            <> "."---- | Type-check a top-level (or module-level) function definition.--- Despite the name, this is also used for checking constant--- definitions, by treating them as 0-ary functions.-checkFunDef ::-  ( VName,-    Maybe (TypeExp (ExpBase NoInfo VName) VName),-    [TypeParam],-    [PatBase NoInfo VName ParamType],-    ExpBase NoInfo VName,-    SrcLoc-  ) ->-  TypeM-    ( [TypeParam],-      [Pat ParamType],-      Maybe (TypeExp Exp VName),-      ResRetType,-      Exp-    )-checkFunDef (fname, maybe_retdecl, tparams, params, body, loc) =-  runTermTypeM checkExp $ do-    (tparams', params', maybe_retdecl', RetType dims rettype', body') <--      checkBinding (fname, maybe_retdecl, tparams, params, body, loc)--    -- Since this is a top-level function, we also resolve overloaded-    -- types, using either defaults or complaining about ambiguities.-    fixOverloadedTypes $-      typeVars rettype' <> foldMap (typeVars . patternType) params'--    -- Then replace all inferred types in the body and parameters.-    body'' <- normTypeFully body'-    params'' <- mapM normTypeFully params'-    maybe_retdecl'' <- traverse updateTypes maybe_retdecl'-    rettype'' <- normTypeFully rettype'--    -- Check if the function body can actually be evaluated.-    causalityCheck body''--    -- Check for various problems.-    mapM_ (mustBeIrrefutable . fmap toStruct) params'-    localChecks body''--    let ((body''', updated_ret), errors) =-          Consumption.checkValDef-            ( fname,-              params'',-              body'',-              RetType dims rettype'',-              maybe_retdecl'',-              loc-            )--    mapM_ throwError errors--    pure (tparams', params'', maybe_retdecl'', updated_ret, body''')---- | This is "fixing" as in "setting them", not "correcting them".  We--- only make very conservative fixing.-fixOverloadedTypes :: Names -> TermTypeM ()-fixOverloadedTypes tyvars_at_toplevel =-  getConstraints >>= mapM_ fixOverloaded . M.toList . M.map snd-  where-    fixOverloaded (v, Overloaded ots usage)-      | Signed Int32 `elem` ots = do-          unify usage (Scalar (TypeVar mempty (qualName v) [])) $-            Scalar (Prim $ Signed Int32)-          when (v `S.member` tyvars_at_toplevel) $-            warn usage "Defaulting ambiguous type to i32."-      | FloatType Float64 `elem` ots = do-          unify usage (Scalar (TypeVar mempty (qualName v) [])) $-            Scalar (Prim $ FloatType Float64)-          when (v `S.member` tyvars_at_toplevel) $-            warn usage "Defaulting ambiguous type to f64."-      | otherwise =-          typeError usage mempty . withIndexLink "ambiguous-type" $-            "Type is ambiguous (could be one of"-              <+> commasep (map pretty ots)-              <> ")."-                </> "Add a type annotation to disambiguate the type."-    fixOverloaded (v, NoConstraint _ usage) = do-      -- See #1552.-      unify usage (Scalar (TypeVar mempty (qualName v) [])) $-        Scalar (tupleRecord [])-      when (v `S.member` tyvars_at_toplevel) $-        warn usage "Defaulting ambiguous type to ()."-    fixOverloaded (_, Equality usage) =-      typeError usage mempty . withIndexLink "ambiguous-type" $-        "Type is ambiguous (must be equality type)."-          </> "Add a type annotation to disambiguate the type."-    fixOverloaded (_, HasFields _ fs usage) =-      typeError usage mempty . withIndexLink "ambiguous-type" $-        "Type is ambiguous.  Must be record with fields:"-          </> indent 2 (stack $ map field $ M.toList fs)-          </> "Add a type annotation to disambiguate the type."-      where-        field (l, t) = pretty l <> colon <+> align (pretty t)-    fixOverloaded (_, HasConstrs _ cs usage) =-      typeError usage mempty . withIndexLink "ambiguous-type" $-        "Type is ambiguous (must be a sum type with constructors:"-          <+> pretty (Sum cs)-          <> ")."-            </> "Add a type annotation to disambiguate the type."-    fixOverloaded (v, Size Nothing (Usage Nothing loc)) =-      typeError loc mempty . withIndexLink "ambiguous-size" $-        "Ambiguous size" <+> dquotes (prettyName v) <> "."-    fixOverloaded (v, Size Nothing (Usage (Just u) loc)) =-      typeError loc mempty . withIndexLink "ambiguous-size" $-        "Ambiguous size" <+> dquotes (prettyName v) <+> "arising from" <+> pretty u <> "."-    fixOverloaded _ = pure ()--hiddenParamNames :: [Pat ParamType] -> [VName]-hiddenParamNames params = hidden-  where-    param_all_names = mconcat $ map patNames params-    named (Named x, _, _) = Just x-    named (Unnamed, _, _) = Nothing-    param_names =-      S.fromList $ mapMaybe (named . patternParam) params-    hidden = filter (`notElem` param_names) param_all_names--inferredReturnType :: SrcLoc -> [Pat ParamType] -> StructType -> TermTypeM StructType-inferredReturnType loc params t = do-  -- The inferred type may refer to names that are bound by the-  -- parameter patterns, but which will not be visible in the type.-  -- These we must turn into fresh type variables, which will be-  -- existential in the return type.-  fst <$> unscopeType loc hidden_params t-  where-    hidden_params = filter (`elem` hidden) $ foldMap patNames params-    hidden = hiddenParamNames params--checkBinding ::-  ( VName,-    Maybe (TypeExp (ExpBase NoInfo VName) VName),-    [TypeParam],-    [PatBase NoInfo VName ParamType],-    ExpBase NoInfo VName,-    SrcLoc-  ) ->-  TermTypeM-    ( [TypeParam],-      [Pat ParamType],-      Maybe (TypeExp Exp VName),-      ResRetType,-      Exp-    )-checkBinding (fname, maybe_retdecl, tparams, params, body, loc) =-  incLevel . bindingParams tparams params $ \params' -> do-    maybe_retdecl' <- traverse checkTypeExpNonrigid maybe_retdecl--    body' <--      checkFunBody-        params'-        body-        ((\(_, x, _) -> x) <$> maybe_retdecl')-        (maybe loc srclocOf maybe_retdecl)--    params'' <- mapM updateTypes params'-    body_t <- expTypeFully body'--    (maybe_retdecl'', rettype) <- case maybe_retdecl' of-      Just (retdecl', ret, _) -> do-        ret' <- normTypeFully ret-        pure (Just retdecl', ret')-      Nothing-        | null params ->-            pure (Nothing, toRes Nonunique body_t)-        | otherwise -> do-            body_t' <- inferredReturnType loc params'' body_t-            pure (Nothing, toRes Nonunique body_t')--    verifyFunctionParams (Just fname) params''--    (tparams', params''', rettype') <--      letGeneralise (baseName fname) loc tparams params'' =<< unscopeUnknown rettype--    when-      ( null params-          && any isSizeParam tparams'-          && not (null (retDims rettype'))-      )-      $ typeError loc mempty-      $ textwrap "A size-polymorphic value binding may not have a type with an existential size."-        </> "Type of this binding is:"-        </> indent 2 (pretty rettype')-        </> "with the following type parameters:"-        </> indent 2 (sep $ map pretty $ filter isSizeParam tparams')--    pure (tparams', params''', maybe_retdecl'', rettype', body')---- | 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-  where-    onParam :: TypeBase Size als -> S.Set VName-    onParam (Scalar Arrow {}) = mempty-    onParam (Scalar (Record fs)) = mconcat $ map onParam $ M.elems fs-    onParam (Scalar (TypeVar _ _ targs)) = mconcat $ map onTypeArg targs-    onParam t = fvVars $ freeInType t-    onTypeArg (TypeArgDim (Var d _ _)) = S.singleton $ qualLeaf d-    onTypeArg (TypeArgDim _) = mempty-    onTypeArg (TypeArgType t) = onParam t-sizeNamesPos _ = mempty---- | Verify certain restrictions on function parameters, and bail out--- on dubious constructions.------ These restrictions apply to all functions (anonymous or otherwise).--- Top-level functions have further restrictions that are checked--- during let-generalisation.-verifyFunctionParams :: Maybe VName -> [Pat ParamType] -> TermTypeM ()-verifyFunctionParams fname params =-  onFailure (CheckingParams (baseName <$> fname)) $-    verifyParams (foldMap patNames params) =<< mapM updateTypes params-  where-    verifyParams forbidden (p : ps)-      | d : _ <- filter (`elem` forbidden) $ S.toList $ fvVars $ freeInPat p =-          typeError p mempty . withIndexLink "inaccessible-size" $-            "Parameter"-              <+> dquotes (pretty p)-              </> "refers to size"-              <+> dquotes (prettyName d)-              <> comma-                </> textwrap "which will not be accessible to the caller"-              <> comma-                </> textwrap "possibly because it is nested in a tuple or record."-                </> textwrap "Consider ascribing an explicit type that does not reference "-              <> dquotes (prettyName d)-              <> "."-      | otherwise = verifyParams forbidden' ps-      where-        forbidden' =-          case patternParam p of-            (Named v, _, _) -> delete v forbidden-            _ -> forbidden-    verifyParams _ [] = pure ()---- | Move existentials down to the level where they are actually used--- (i.e. have their "witnesses").  E.g. changes------ @--- ?[n].bool -> [n]bool--- @------ to------ @--- bool -> ?[n].[n]bool--- @-injectExt :: [VName] -> TypeBase Size u -> RetTypeBase Size u-injectExt [] ret = RetType [] ret-injectExt ext ret = RetType ext_here $ deeper ret-  where-    (immediate, _) = dimUses ret-    (ext_here, ext_there) = partition (`S.member` immediate) ext-    deeper :: TypeBase Size u -> TypeBase Size u-    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 d1 t1 (RetType t2_ext t2))) =-      Scalar $ Arrow als p d1 t1 $ injectExt (nubOrd (ext_there <> t2_ext)) t2-    deeper (Scalar (TypeVar u tn targs)) =-      Scalar $ TypeVar u tn $ map deeperArg targs-    deeper t@Array {} = t--    deeperArg (TypeArgType t) = TypeArgType $ deeper t-    deeperArg (TypeArgDim d) = TypeArgDim d---- | Find all type variables in the given type that are covered by the--- constraints, and produce type parameters that close over them.------ The passed-in list of type parameters is always prepended to the--- produced list of type parameters.-closeOverTypes ::-  Name ->-  SrcLoc ->-  [TypeParam] ->-  [StructType] ->-  ResType ->-  Constraints ->-  TermTypeM ([TypeParam], ResRetType)-closeOverTypes defname defloc tparams paramts ret substs = do-  (more_tparams, retext) <--    partitionEithers . catMaybes-      <$> mapM closeOver (M.toList $ M.map snd to_close_over)-  let mkExt v =-        case M.lookup v substs of-          Just (_, UnknownSize {}) -> Just v-          _ -> Nothing-  pure-    ( tparams ++ more_tparams,-      injectExt (nubOrd $ retext ++ mapMaybe mkExt (S.toList $ fvVars $ freeInType ret)) ret-    )-  where-    -- Diet does not matter here.-    t = foldFunType (map (toParam Observe) paramts) $ RetType [] ret-    to_close_over = M.filterWithKey (\k _ -> k `S.member` visible) substs-    visible = typeVars t <> fvVars (freeInType t)--    (produced_sizes, param_sizes) = dimUses t--    -- Avoid duplicate type parameters.-    closeOver (k, _)-      | k `elem` map typeParamName tparams =-          pure Nothing-    closeOver (k, NoConstraint l _) =-      pure $ Just $ Left $ TypeParamType l k mempty-    closeOver (k, ParamType l _) =-      pure $ Just $ Left $ TypeParamType l k mempty-    closeOver (k, Size Nothing _) =-      pure $ Just $ Left $ TypeParamDim k mempty-    closeOver (k, UnknownSize _ _)-      | k `S.member` param_sizes,-        k `S.notMember` produced_sizes = do-          notes <- dimNotes defloc $ sizeFromName (qualName k) mempty-          typeError defloc notes . withIndexLink "unknown-param-def" $-            "Unknown size"-              <+> dquotes (prettyName k)-              <+> "in parameter of"-              <+> dquotes (prettyName defname)-              <> ", which is inferred as:"-                </> indent 2 (pretty t)-      | k `S.member` produced_sizes =-          pure $ Just $ Right k-    closeOver (_, _) =-      pure Nothing--letGeneralise ::-  Name ->-  SrcLoc ->-  [TypeParam] ->-  [Pat ParamType] ->-  ResType ->-  TermTypeM ([TypeParam], [Pat ParamType], ResRetType)-letGeneralise defname defloc tparams params restype =-  onFailure (CheckingLetGeneralise defname) $ do-    now_substs <- getConstraints--    -- Candidates for let-generalisation are those type variables that-    ---    -- (1) were not known before we checked this function, and-    ---    -- (2) are not used in the (new) definition of any type variables-    -- known before we checked this function.-    ---    -- (3) are not referenced from an overloaded type (for example,-    -- are the element types of an incompletely resolved record type).-    -- This is a bit more restrictive than I'd like, and SML for-    -- example does not have this restriction.-    ---    -- Criteria (1) and (2) is implemented by looking at the binding-    -- level of the type variables.-    let keep_type_vars = overloadedTypeVars now_substs--    cur_lvl <- curLevel-    let candidate k (lvl, _) = (k `S.notMember` keep_type_vars) && lvl >= (cur_lvl - length params)-        new_substs = M.filterWithKey candidate now_substs--    (tparams', RetType ret_dims restype') <--      closeOverTypes-        defname-        defloc-        tparams-        (map patternStructType params)-        restype-        new_substs--    restype'' <- updateTypes restype'--    let used_sizes =-          freeInType restype'' <> foldMap (freeInType . patternType) params-    case filter ((`S.notMember` fvVars used_sizes) . typeParamName) $-      filter isSizeParam tparams' of-      [] -> pure ()-      tp : _ -> unusedSize $ SizeBinder (typeParamName tp) (srclocOf tp)--    -- We keep those type variables that were not closed over by-    -- let-generalisation.-    modifyConstraints $ M.filterWithKey $ \k _ -> k `notElem` map typeParamName tparams'--    pure (tparams', params, RetType ret_dims restype'')--checkFunBody ::-  [Pat ParamType] ->-  ExpBase NoInfo VName ->-  Maybe ResType ->-  SrcLoc ->-  TermTypeM Exp-checkFunBody params body maybe_rettype loc = do-  body' <- checkExp body--  -- Unify body return type with return annotation, if one exists.-  case maybe_rettype of-    Just rettype -> do-      body_t <- expTypeFully body'-      -- We need to turn any sizes provided by "hidden" parameter-      -- names into existential sizes instead.-      let hidden = hiddenParamNames params-      (body_t', _) <--        unscopeType-          loc-          (filter (`elem` hidden) $ foldMap patNames params)-          body_t-      case find (`elem` hidden) $ fvVars $ freeInType rettype of-        Just v ->-          typeError loc mempty $-            "The return type annotation"-              </> indent 2 (align (pretty rettype))-              </> "refers to the name"-              <+> dquotes (prettyName v)-              <+> "which is bound to an inner component of a function parameter."-        Nothing -> do-          let usage = mkUsage body "return type annotation"-          onFailure (CheckingReturn rettype body_t') $-            unify usage (toStruct rettype) body_t'-    Nothing -> pure ()--  pure body'--arrayOfM ::-  SrcLoc ->-  StructType ->-  Shape Size ->-  TermTypeM StructType-arrayOfM loc t shape = do-  arrayElemType (mkUsage loc "use as array element") "type used in array" t-  pure $ arrayOf shape t+--+-- The strategy is to split type checking into sveral (main) passes:+--+-- 1) A size-agnostic pass implemented in+-- "Language.Futhark.TypeChecker.Terms.Unsized".+--+-- 2) Pass (1) has given us a program where we know the types of+-- everything, but the sizes of nothing. Pass (2) then does+-- essentially size inference, with the benefit of already knowing the+-- full unsized type of everything. This is done using a syntax-driven+-- approach, similar to Algorithm W.+--+-- 3) The program is then checked for violation of uniqueness+-- properties, which is implemented in+-- "Language.Futhark.TypeChecker.Consumption".+module Language.Futhark.TypeChecker.Terms+  ( checkOneExp,+    checkSizeExp,+    checkFunDef,+  )+where++import Control.Monad+import Control.Monad.Except+import Control.Monad.Identity+import Control.Monad.Reader+import Control.Monad.State.Strict+import Data.Bifunctor+import Data.Bitraversable+import Data.Char (isAscii)+import Data.Either+import Data.List (delete, find, genericLength, partition)+import Data.List qualified as L+import Data.List.NonEmpty qualified as NE+import Data.Map.Strict qualified as M+import Data.Maybe+import Data.Set qualified as S+import Data.Text qualified as T+import Futhark.Util (mapAccumLM, nubOrd, topologicalSort)+import Futhark.Util.Pretty hiding (space)+import Language.Futhark+import Language.Futhark.Primitive (intByteSize)+import Language.Futhark.Traversals+import Language.Futhark.TypeChecker.Consumption qualified as Consumption+import Language.Futhark.TypeChecker.Match+import Language.Futhark.TypeChecker.Monad hiding (BoundV, lookupAbsTy, lookupMod)+import Language.Futhark.TypeChecker.Terms.Loop+import Language.Futhark.TypeChecker.Terms.Monad+import Language.Futhark.TypeChecker.Terms.Pat+import Language.Futhark.TypeChecker.Terms.Unsized qualified as Unsized+import Language.Futhark.TypeChecker.Types+import Language.Futhark.TypeChecker.Unify+import Prelude hiding (mod)++hasBinding :: Exp -> Bool+hasBinding Lambda {} = True+hasBinding (AppExp LetPat {} _) = True+hasBinding (AppExp LetFun {} _) = True+hasBinding (AppExp Loop {} _) = True+hasBinding (AppExp LetWith {} _) = True+hasBinding (AppExp Match {} _) = True+hasBinding e = isNothing $ astMap m e+  where+    m =+      identityMapper {mapOnExp = \e' -> if hasBinding e' then Nothing else Just e'}++--- Basic checking++-- | Determine if the two types are identical, ignoring uniqueness.+-- Mismatched dimensions are turned into fresh rigid type variables.+-- Causes a 'TypeError' if they fail to match, and otherwise returns+-- one of them.+unifyBranchTypes :: SrcLoc -> StructType -> StructType -> TermTypeM (StructType, [VName])+unifyBranchTypes loc t1 t2 =+  onFailure (CheckingBranches t1 t2) $+    unifyMostCommon (mkUsage loc "unification of branch results") t1 t2++unifyBranches :: SrcLoc -> Exp -> Exp -> TermTypeM (StructType, [VName])+unifyBranches loc e1 e2 = do+  e1_t <- expType e1+  e2_t <- expType e2+  unifyBranchTypes loc e1_t e2_t++sliceShape ::+  Maybe (SrcLoc, Rigidity) ->+  [DimIndex] ->+  TypeBase Size as ->+  TermTypeM (TypeBase Size as, [VName])+sliceShape r slice t@(Array u (Shape orig_dims) et) =+  runStateT (setDims <$> adjustDims slice orig_dims) []+  where+    setDims [] = stripArray (length orig_dims) t+    setDims dims' = Array u (Shape dims') et++    -- If the result is supposed to be a nonrigid size variable, then+    -- don't bother trying to create non-existential sizes.  This is+    -- necessary to make programs type-check without too much+    -- ceremony; see e.g. tests/inplace5.fut.+    isRigid Rigid {} = True+    isRigid _ = False+    refine_sizes = maybe False (isRigid . snd) r++    sliceSize orig_d i j stride =+      case r of+        Just (loc, Rigid _) -> do+          (d, ext) <-+            lift . extSize loc $+              SourceSlice orig_d' (bareExp <$> i) (bareExp <$> j) (bareExp <$> stride)+          modify (maybeToList ext ++)+          pure d+        Just (loc, Nonrigid) ->+          lift $+            flip sizeFromName loc . qualName+              <$> newFlexibleDim (mkUsage loc "size of slice") "slice_dim"+        Nothing -> do+          v <- lift $ newID "slice_anydim"+          modify (v :)+          pure $ sizeFromName (qualName v) mempty+      where+        -- The original size does not matter if the slice is fully specified.+        orig_d'+          | isJust i, isJust j = Nothing+          | otherwise = Just orig_d++    warnIfBinding binds d i j stride size =+      if binds+        then do+          lift . warn (srclocOf size) $+            withIndexLink+              "size-expression-bind"+              "Size expression with binding is replaced by unknown size."+          (:) <$> sliceSize d i j stride+        else pure (size :)++    adjustDims (DimFix {} : idxes') (_ : dims) =+      adjustDims idxes' dims+    -- Pat match some known slices to be non-existential.+    adjustDims (DimSlice i j stride : idxes') (d : dims)+      | refine_sizes,+        maybe True ((== Just 0) . isInt64) i,+        maybe True ((== Just 1) . isInt64) stride = do+          let binds = maybe False hasBinding j+          warnIfBinding binds d i j stride (fromMaybe d j)+            <*> adjustDims idxes' dims+    adjustDims ((DimSlice i j stride) : idxes') (d : dims)+      | refine_sizes,+        Just i' <- i, -- if i ~ 0, previous case+        maybe True ((== Just 1) . isInt64) stride = do+          let j' = fromMaybe d j+              binds = hasBinding j' || hasBinding i'+          warnIfBinding binds d i j stride (sizeMinus j' i')+            <*> adjustDims idxes' dims+    -- stride == -1+    adjustDims ((DimSlice Nothing Nothing stride) : idxes') (d : dims)+      | refine_sizes,+        maybe True ((== Just (-1)) . isInt64) stride =+          (d :) <$> adjustDims idxes' dims+    adjustDims ((DimSlice (Just i) (Just j) stride) : idxes') (d : dims)+      | refine_sizes,+        maybe True ((== Just (-1)) . isInt64) stride = do+          let binds = hasBinding i || hasBinding j+          warnIfBinding binds d (Just i) (Just j) stride (sizeMinus i j)+            <*> adjustDims idxes' dims+    -- existential+    adjustDims ((DimSlice i j stride) : idxes') (d : dims) =+      (:) <$> sliceSize d i j stride <*> adjustDims idxes' dims+    adjustDims _ dims =+      pure dims++    sizeMinus j i =+      AppExp+        ( BinOp+            (qualName (intrinsicVar "-"), mempty)+            sizeBinOpInfo+            (j, Info Nothing)+            (i, Info Nothing)+            mempty+        )+        $ Info+        $ AppRes i64 []+    i64 = Scalar $ Prim $ Signed Int64+    sizeBinOpInfo = Info $ foldFunType [i64, i64] $ RetType [] i64+sliceShape _ _ t = pure (t, [])++--- Main checkers++checkAscript ::+  SrcLoc ->+  TypeExp Exp VName ->+  Exp ->+  TermTypeM (TypeExp Exp VName, Exp)+checkAscript loc te e = do+  (te', decl_t, _) <- checkTypeExpNonrigid te+  e' <- checkExp e+  e_t <- expTypeFully e'++  onFailure (CheckingAscription (toStruct decl_t) e_t) $+    unify (mkUsage loc "type ascription") (toStruct decl_t) e_t++  pure (te', e')++checkCoerce ::+  SrcLoc ->+  TypeExp Exp VName ->+  Exp ->+  TermTypeM (TypeExp Exp VName, StructType, Exp)+checkCoerce loc te e = do+  (te', te_t, ext) <- checkTypeExpNonrigid te+  e' <- checkExp e+  e_t <- expTypeFully e'++  te_t_nonrigid <- makeNonExtFresh ext $ toStruct te_t++  onFailure (CheckingAscription (toStruct te_t) e_t) $+    unify (mkUsage loc "size coercion") e_t te_t_nonrigid++  -- If the type expression had any anonymous dimensions, these will+  -- now be in 'ext'.  Those we keep nonrigid and unify with e_t.+  -- This ensures that 'x :> [1][]i32' does not make the second+  -- dimension unknown.  Use of matchDims is sensible because the+  -- structure of e_t' will be fully known due to the unification, and+  -- te_t because type expressions are complete.+  pure (te', toStruct te_t, e')+  where+    makeNonExtFresh ext = bitraverse onDim pure+      where+        onDim d@(Var v _ _)+          | qualLeaf v `elem` ext = pure d+        onDim d = do+          v <- newTypeName "coerce"+          constrain v . Size Nothing $+            mkUsage+              loc+              "a size coercion where the underlying expression size cannot be determined"+          pure $ sizeFromName (qualName v) (srclocOf d)++-- Expressions witnessed by type, topologically sorted.+topWit :: TypeBase Exp u -> [Exp]+topWit = topologicalSort depends . witnessedExps+  where+    witnessedExps t = execState (traverseDims onDim t) mempty+      where+        onDim _ PosImmediate e = modify (e :)+        onDim _ _ _ = pure ()+    depends a b = any (sameExp b) $ subExps a++sizeFree ::+  (MonadUnify m) =>+  SrcLoc ->+  (Exp -> Maybe VName) ->+  TypeBase Size u ->+  m (TypeBase Size u, [VName])+sizeFree tloc expKiller orig_t = do+  runReaderT (toBeReplaced orig_t $ onType orig_t) mempty `runStateT` mempty+  where+    lookReplacement e repl = snd <$> L.find (sameExp e . fst) repl+    expReplace mapping e+      | Just e' <- lookReplacement e mapping = e'+      | otherwise = runIdentity $ astMap mapper e+      where+        mapper = identityMapper {mapOnExp = pure . expReplace mapping}++    replacing e = do+      e' <- asks (`expReplace` e)+      case expKiller e' of+        Nothing -> pure e'+        Just cause -> do+          vn <- lift $ lift $ newRigidDim tloc (RigidOutOfScope (locOf e) cause) "d"+          modify (vn :)+          pure $ sizeFromName (qualName vn) (srclocOf e)++    toBeReplaced t m' = foldl f m' $ topWit t+      where+        f m e = do+          e' <- replacing e+          local ((e, e') :) m++    onScalar (Record fs) =+      Record <$> traverse onType fs+    onScalar (Sum cs) =+      Sum <$> (traverse . traverse) onType cs+    onScalar (Arrow as pn d argT (RetType dims retT)) = do+      argT' <- onType argT+      old_bound <- get+      retT' <- toBeReplaced retT $ onType retT+      rl <- state $ L.partition (`notElem` old_bound)+      let dims' = dims <> rl+      pure $ Arrow as pn d argT' (RetType dims' retT')+    onScalar (TypeVar u v args) =+      TypeVar u v <$> mapM onTypeArg args+      where+        onTypeArg (TypeArgDim d) = TypeArgDim <$> replacing d+        onTypeArg (TypeArgType ty) = TypeArgType <$> onType ty+    onScalar (Prim pt) = pure $ Prim pt++    onType ::+      (MonadUnify m) =>+      TypeBase Size u ->+      ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Size u)+    onType (Array u shape scalar) =+      Array u <$> traverse replacing shape <*> onScalar scalar+    onType (Scalar ty) =+      Scalar <$> onScalar ty++-- Remove unknown sizes from function body types before we perform+-- let-generalisation. This is because if a function is inferred to return+-- something of type '[x+y]t' where 'x' or 'y' are unknown, we want to turn that+-- into '[z]t', where 'z' is a fresh unknown, which is then by+-- let-generalisation turned into '?[z].[z]t'.+unscopeUnknown ::+  TypeBase Size u ->+  TermTypeM (TypeBase Size u)+unscopeUnknown t = do+  constraints <- getConstraints+  -- The killer only ever fires on an unknown-size variable, so if none occurs+  -- free in the type there is nothing to do and we can skip the traversal (and+  -- the witness computation) entirely.+  if not (any (isUnknown constraints) (fvVars (freeInType t)))+    then pure t+    else -- These sizes will be immediately turned into existentials, so we do+    -- not need to care about their location.+      fst <$> sizeFree mempty (expKiller constraints) t+  where+    expKiller _ Var {} = Nothing+    expKiller constraints e =+      S.lookupMin $ S.filter (isUnknown constraints) $ (`S.difference` witnesses) $ fvVars $ freeInExp e+    isUnknown constraints vn+      | Just UnknownSize {} <- snd <$> M.lookup vn constraints = True+    isUnknown _ _ = False+    (witnesses, _) = determineSizeWitnesses $ toStruct t++unscopeType ::+  SrcLoc ->+  [VName] ->+  TypeBase Size as ->+  TermTypeM (TypeBase Size as, [VName])+unscopeType tloc unscoped t+  -- Fast-path for common case where 't' has no free variables in unscoped.+  | not (any (`elem` unscoped) (fvVars (freeInType t))) = pure (t, [])+  | otherwise =+      sizeFree tloc (find (`elem` unscoped) . fvVars . freeInExp) t++checkExp :: Exp -> TermTypeM Exp+checkExp (Var qn (Info t) loc) = do+  t' <- lookupVar loc qn t+  pure $ Var qn (Info t') loc+checkExp (Literal val loc) =+  pure $ Literal val loc+checkExp (Hole (Info t) loc) = do+  t' <- replaceTyVarsAbsorbable loc t+  pure $ Hole (Info t') loc+checkExp (StringLit vs loc) =+  pure $ StringLit vs loc+checkExp (IntLit val (Info t) loc) = do+  t' <- replaceTyVars loc t+  pure $ IntLit val (Info t') loc+checkExp (FloatLit val (Info t) loc) = do+  t' <- replaceTyVars loc t+  pure $ FloatLit val (Info t') loc+checkExp (TupLit es loc) =+  TupLit <$> mapM checkExp es <*> pure loc+checkExp (RecordLit fs loc) =+  RecordLit <$> mapM checkField fs <*> pure loc+  where+    checkField (RecordFieldExplicit f e rloc) =+      RecordFieldExplicit f <$> checkExp e <*> pure rloc+    checkField (RecordFieldImplicit name (Info t) rloc) = do+      t' <- lookupVar rloc (qualName (unLoc name)) t+      pure $ RecordFieldImplicit name (Info t') rloc+-- No need to type check this, as these are only produced by the+-- parser if the elements are monomorphic and all match.+checkExp (ArrayVal vs t loc) =+  pure $ ArrayVal vs t loc+checkExp (ArrayLit all_es (Info t) loc) =+  -- We only consult the type inferred by the unsized type checker+  -- for empty arrays; otherwise we use the type of the first+  -- element.  This significantly cuts down on the number of+  -- inferred types we have to instantiate for pathologically large+  -- multidimensional array literals.+  case all_es of+    [] -> do+      t' <- replaceTyVars loc t+      case peelArray 1 t' of+        Just et -> do+          let t'' = arrayOf (Shape [sizeFromInteger 0 mempty]) et+          unify (mkUsage loc "empty array literal") t'' t'+          pure $ ArrayLit [] (Info t'') loc+        Nothing -> error $ "checkExp ArrayLit: " <> prettyString t'+    e : es -> do+      e' <- checkExp e+      et <- expType e'+      es' <- mapM (unifies "type of first array element" et <=< checkExp) es+      let arr_t = arrayOf (Shape [sizeFromInteger (genericLength all_es) mempty]) et+      pure $ ArrayLit (e' : es') (Info arr_t) loc+checkExp (AppExp (Range start maybe_step end loc) _) = do+  start' <- checkExp start+  start_t <- expType start'+  maybe_step' <- case maybe_step of+    Nothing -> pure Nothing+    Just step -> do+      let warning = warn loc "First and second element of range are identical, this will produce an empty array."+      case (start, step) of+        (Literal x _, Literal y _) -> when (x == y) warning+        (Var x_name _ _, Var y_name _ _) -> when (x_name == y_name) warning+        _ -> pure ()+      Just <$> (unifies "use in range expression" start_t =<< checkExp step)++  let unifyRange e = unifies "use in range expression" start_t =<< checkExp e+  end' <- traverse unifyRange end++  end_t <- case end' of+    DownToExclusive e -> expType e+    ToInclusive e -> expType e+    UpToExclusive e -> expType e++  -- Special case some ranges to give them a known size.+  let warnIfBinding binds size =+        if binds+          then do+            warn (srclocOf size) $+              withIndexLink+                "size-expression-bind"+                "Size expression with binding is replaced by unknown size."+            d <- newRigidDim loc RigidRange "range_dim"+            pure (sizeFromName (qualName d) mempty, Just d)+          else pure (size, Nothing)+  (dim, retext) <-+    case (isInt64 start', isInt64 <$> maybe_step', end') of+      (Just 0, Just (Just 1), UpToExclusive end'')+        | Scalar (Prim (Signed Int64)) <- end_t ->+            warnIfBinding (hasBinding end'') end''+      (Just 0, Nothing, UpToExclusive end'')+        | Scalar (Prim (Signed Int64)) <- end_t ->+            warnIfBinding (hasBinding end'') end''+      (_, Nothing, UpToExclusive end'')+        | Scalar (Prim (Signed Int64)) <- end_t ->+            warnIfBinding (hasBinding end'' || hasBinding start') $ sizeMinus end'' start'+      (_, Nothing, ToInclusive end'')+        -- No stride means we assume a stride of one.+        | Scalar (Prim (Signed Int64)) <- end_t ->+            warnIfBinding (hasBinding end'' || hasBinding start') $ sizeMinusInc end'' start'+      (Just 1, Just (Just 2), ToInclusive end'')+        | Scalar (Prim (Signed Int64)) <- end_t ->+            warnIfBinding (hasBinding end'') end''+      _ -> do+        d <- newRigidDim loc RigidRange "range_dim"+        pure (sizeFromName (qualName d) mempty, Just d)++  let t = arrayOf (Shape [dim]) start_t+      res = AppRes t (maybeToList retext)++  pure $ AppExp (Range start' maybe_step' end' loc) (Info res)+  where+    i64 = Scalar $ Prim $ Signed Int64+    mkBinOp op t x y =+      AppExp+        ( BinOp+            (qualName (intrinsicVar op), mempty)+            sizeBinOpInfo+            (x, Info Nothing)+            (y, Info Nothing)+            mempty+        )+        (Info $ AppRes t [])+    mkSub = mkBinOp "-" i64+    mkAdd = mkBinOp "+" i64+    sizeMinus j i = j `mkSub` i+    sizeMinusInc j i = (j `mkSub` i) `mkAdd` sizeFromInteger 1 mempty+    sizeBinOpInfo = Info $ foldFunType [i64, i64] $ RetType [] i64+checkExp (Ascript e te loc) = do+  (te', e') <- checkAscript loc te e+  pure $ Ascript e' te' loc+checkExp (Coerce e te _ loc) = do+  (te', te_t, e') <- checkCoerce loc te e+  t <- expTypeFully e'+  t' <- matchDims (const . const pure) t te_t+  pure $ Coerce e' te' (Info t') loc+checkExp (AppExp (Apply fe args loc) _) = do+  fe' <- checkExp fe+  args' <- mapM (checkExp . 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+      (_, rt, argext, exts) <- checkApply loc (fname, i) t arg'+      pure+        ( (i + 1, all_exts <> exts, rt),+          (Info argext, arg')+        )+checkExp (AppExp (BinOp (op, oploc) (Info op_t) (e1, _) (e2, _) loc) _) = do+  ftype <- lookupVar oploc op op_t+  e1' <- checkExp e1+  e2' <- checkExp e2+  -- Note that the application to the first operand cannot fix any+  -- existential sizes, because it must by necessity be a function.+  (_, rt, p1_ext, _) <- checkApply loc (Just op, 0) ftype e1'+  (_, rt', p2_ext, retext) <- checkApply loc (Just op, 1) rt e2'++  pure $+    AppExp+      ( BinOp+          (op, oploc)+          (Info ftype)+          (e1', Info p1_ext)+          (e2', Info p2_ext)+          loc+      )+      (Info (AppRes rt' retext))+checkExp (Project k e _ loc) = do+  e' <- checkExp e+  t <- expType e'+  case t of+    Scalar (Record fs)+      | Just kt <- M.lookup k fs ->+          pure $ Project k e' (Info kt) loc+    _ -> error $ "checkExp Project: " <> show t+checkExp (Parens e loc) =+  Parens <$> checkExp e <*> pure loc+checkExp (QualParens (modname, modnameloc) e loc) = do+  mod <- lookupMod modname+  case mod of+    ModEnv env -> local (`withEnv` env) $ do+      e' <- checkExp e+      pure $ QualParens (modname, modnameloc) e' loc+    ModFun {} ->+      typeError loc mempty . withIndexLink "module-is-parametric" $+        "Module" <+> pretty modname <+> " is a parametric module."+checkExp (Negate arg loc) = do+  arg' <- checkExp arg+  pure $ Negate arg' loc+checkExp (Not arg loc) = do+  arg' <- checkExp arg+  pure $ Not arg' loc+checkExp (AppExp (LetPat sizes pat e body loc) _) = do+  e' <- checkExp e++  -- Not technically an ascription, but we want the pattern to have+  -- exactly the type of 'e'.+  t <- expType e'+  bindingSizes sizes . incLevel . bindingPat sizes pat t $ \pat' -> do+    body' <- incLevel $ checkExp body+    body_t <- expTypeFully body'++    -- If the bound expression is of type i64, then we replace the+    -- pattern name with the expression in the type of the body.+    -- Otherwise, we need to come up with unknown sizes for the+    -- sizes going out of scope.+    (body_t', retext) <-+      case (t, patNames pat') of+        (Scalar (Prim (Signed Int64)), [v])+          | not $ hasBinding e' -> do+              let f x = if x == v then Just (ExpSubst e') else Nothing+              pure (applySubst f body_t, [])+        _ ->+          unscopeType loc (map sizeName sizes <> patNames pat') body_t++    pure $+      AppExp+        (LetPat sizes (fmap toStruct pat') e' body' loc)+        (Info $ AppRes body_t' retext)+checkExp (AppExp (LetFun name (tparams, params, maybe_retdecl, _, e) body loc) _) = do+  (tparams', params', maybe_retdecl', rettype, e') <-+    checkBinding (fst name, maybe_retdecl, tparams, params, e, loc)++  let entry = BoundV tparams' $ funType params' rettype+      bindF scope =+        scope+          { scopeVtable = M.insert (fst name) entry $ scopeVtable scope+          }+  body' <- localScope bindF $ checkExp body++  (body_t, ext) <- unscopeType loc [fst name] =<< expTypeFully body'++  pure $+    AppExp+      ( LetFun+          name+          (tparams', params', maybe_retdecl', Info rettype, e')+          body'+          loc+      )+      (Info $ AppRes body_t ext)+checkExp (AppExp (LetWith dest src steps ve body loc) _) = do+  -- The type recorded in the AST is the unsized type from the+  -- unsized type checker; we must consult the scope to get the+  -- actual type of the source variable.+  src_t <-+    normTypeFully+      =<< lookupVar (srclocOf src) (qualName $ identName src) (unInfo $ identType src)+  let src' = src {identType = Info src_t}++  case mapAndUnzipM isField steps of+    Just (steps', names) -> do+      ve' <- checkExp ve+      ve_t <- expType ve'+      updated_t <- updateFieldPath src names ve_t src_t++      let dest' = dest {identType = Info updated_t}+      bindingIdent dest' $ do+        body' <- checkExp body+        (body_t, ext) <- unscopeType loc [identName dest'] =<< expTypeFully body'+        pure $ AppExp (LetWith dest' src' steps' ve' body' loc) (Info $ AppRes body_t ext)+    Nothing -> do+      (steps', target_t) <- checkUpdateSteps loc src_t steps+      ve' <- unifies "type of update target" target_t =<< checkExp ve++      let dest' = dest {identType = Info src_t}+      bindingIdent dest' $ do+        body' <- checkExp body+        (body_t, ext) <- unscopeType loc [identName dest'] =<< expTypeFully body'+        pure $ AppExp (LetWith dest' src' steps' ve' body' loc) (Info $ AppRes body_t ext)+  where+    isField (UpdateStepField f) = Just (UpdateStepField f, f)+    isField _ = Nothing++-- Record updates are a bit hacky, because we do not have row typing+-- (yet?).  For now, we only permit record updates where we know the+-- full type up to the field we are updating.+checkExp (Update src steps ve _ loc) = do+  src' <- checkExp src+  src_t <- expTypeFully src'+  case mapAndUnzipM isField steps of+    Just (steps', names) -> do+      ve' <- checkExp ve+      ve_t <- expType ve'+      updated_t <- updateFieldPath src names ve_t src_t+      pure $ Update src' steps' ve' (Info updated_t) loc+    Nothing -> do+      (steps', target_t) <- checkUpdateSteps loc src_t steps+      ve' <- unifies "type of update target" target_t =<< checkExp ve+      src_t' <- expTypeFully src'+      pure $ Update src' steps' ve' (Info src_t') loc+  where+    isField (UpdateStepField f) = Just (UpdateStepField f, f)+    isField _ = Nothing+checkExp (AppExp (Index e slice loc) _) = do+  slice' <- checkSlice slice+  e' <- checkExp e+  -- XXX, the RigidSlice here will be overridden in sliceShape with a proper value.+  (t', retext) <-+    sliceShape (Just (loc, Rigid (RigidSlice Nothing ""))) slice'+      =<< expTypeFully e'++  pure $ AppExp (Index e' slice' loc) (Info $ AppRes t' retext)+checkExp (Assert e1 e2 _ loc) = do+  e1' <- checkExp e1+  e2' <- checkExp e2+  pure $ Assert e1' e2' (Info (prettyText e1)) loc+checkExp (Lambda params body rettype_te (Info (RetType _ rt)) loc) = do+  (params', body', rettype', RetType dims ty) <-+    incLevel . bindingParams [] params $ \params' -> do+      -- The sizes of the return type are absorbable, as a lambda+      -- returns whatever type the context requires. See Note [Size+      -- Inference].+      rt' <- replaceTyVarsAbsorbable loc rt+      rettype_checked <- traverse checkTypeExpNonrigid rettype_te+      declared_rettype <-+        case rettype_checked of+          Just (_, st, _) -> do+            unify (mkUsage body "lambda return type ascription") (toStruct rt') (toStruct st)+            pure $ Just st+          Nothing -> pure Nothing+      (body', body_t) <- checkFunBody params' body declared_rettype loc++      unify (mkUsage body "inferred return type") (toStruct rt') body_t++      params'' <- mapM updateTypes params'++      -- A lambda has no let-generalisation to decide where its+      -- existential sizes go, so we infer them here - also for a+      -- declared return type, whose quantified sizes may by now have+      -- been solved to unknown sizes.+      rettype_st <-+        inferReturnSizes params'' =<< case rettype_checked of+          Just (_, ret, _) -> normTypeFully ret+          Nothing -> pure $ toRes Nonunique body_t++      pure (params'', body', (\(te, _, _) -> te) <$> rettype_checked, rettype_st)++  verifyFunctionParams Nothing params'++  (ty', dims') <- unscopeType loc dims ty++  pure $ Lambda params' body' rettype' (Info (RetType dims' ty')) loc+  where+    -- Inferring the sizes of the return type of a lambda is a lot+    -- like let-generalisation.  We wish to remove any rigid sizes+    -- that were created when checking the body, except for those that+    -- are visible in types that existed before we entered the body,+    -- are parameters, or are used in parameters.+    inferReturnSizes params' ret = do+      cur_lvl <- curLevel+      let named (Named x, _, _) = Just x+          named (Unnamed, _, _) = Nothing+          param_names = mapMaybe (named . patternParam) params'+          pos_sizes =+            sizeNamesPos $ funType params' $ RetType [] ret+          -- Only rigid sizes computed by the body can be hidden. A+          -- size that is still flexible has not been determined yet,+          -- and hiding it would sever its connection to whatever the+          -- enclosing context determines it to be.+          rigid UnknownSize {} = True+          rigid _ = False+          hide k (lvl, c) =+            rigid c && lvl >= cur_lvl && k `notElem` param_names && k `S.notMember` pos_sizes++      hidden_sizes <-+        S.fromList . M.keys . M.filterWithKey hide <$> getConstraints++      let onDim name+            | name `S.member` hidden_sizes = S.singleton name+          onDim _ = mempty++      pure $ RetType (S.toList $ foldMap onDim $ fvVars $ freeInType ret) ret+checkExp (OpSection op (Info op_t) loc) = do+  ftype <- lookupVar loc op op_t+  pure $ OpSection op (Info ftype) loc+checkExp (OpSectionLeft op (Info op_t) e _ _ loc) = do+  ftype <- lookupVar loc op op_t+  e' <- checkExp e+  (t1, rt, argext, retext) <- checkApply loc (Just op, 0) ftype e'+  case (ftype, rt) of+    (Scalar (Arrow _ m1 d1 _ _), Scalar (Arrow _ m2 d2 t2 (RetType ds rt2))) ->+      pure $+        OpSectionLeft+          op+          (Info ftype)+          e'+          (Info (m1, toParam d1 t1, argext), Info (m2, toParam d2 t2))+          (Info $ RetType ds rt2, Info retext)+          loc+    _ ->+      typeError loc mempty $+        "Operator section with invalid operator of type" <+> pretty ftype+checkExp (OpSectionRight op (Info op_t) e _ _ loc) = do+  ftype <- lookupVar loc op op_t+  e' <- checkExp e+  case ftype of+    Scalar (Arrow _ m1 d1 t1 (RetType [] (Scalar (Arrow _ m2 d2 t2 (RetType dims2 ret))))) -> do+      (t2', arrow', argext, _) <-+        checkApply+          loc+          (Just op, 1)+          (Scalar $ Arrow mempty m2 d2 t2 $ RetType [] $ Scalar $ Arrow Nonunique m1 d1 t1 $ RetType dims2 ret)+          e'+      case arrow' of+        Scalar (Arrow _ _ _ t1' (RetType dims2' ret')) ->+          pure $+            OpSectionRight+              op+              (Info ftype)+              e'+              (Info (m1, toParam d1 t1'), Info (m2, toParam d2 t2', argext))+              (Info $ RetType dims2' ret')+              loc+        _ -> error $ "OpSectionRight: impossible type\n" <> prettyString arrow'+    _ ->+      typeError loc mempty $+        "Operator section with invalid operator of type" <+> pretty ftype+checkExp (UpdateSection steps (Info ft) loc) = do+  -- The unsized type checker has already determined the type of the+  -- parameter; we just have to instantiate its sizes. The result+  -- type is then computed by walking the steps, such that its sizes+  -- are those of the corresponding components of the parameter type.+  a <- case ft of+    Scalar (Arrow _ _ _ pt _) -> replaceTyVars loc pt+    _ -> error $ "checkExp UpdateSection: " <> prettyString ft+  (steps', b, retext) <- checkSectionSteps a steps+  let ft' = Scalar $ Arrow mempty Unnamed Observe a $ RetType retext $ toRes Nonunique b+  pure $ UpdateSection steps' (Info ft') loc+  where+    checkSectionSteps t [] =+      pure ([], t, [])+    checkSectionSteps t (step : rest) =+      case step of+        UpdateStepField f -> do+          t' <- normTypeFully t+          case t' of+            Scalar (Record fs)+              | Just f_t <- M.lookup f fs -> do+                  (rest', target_t, retext) <- checkSectionSteps f_t rest+                  pure (UpdateStepField f : rest', target_t, retext)+            _ ->+              error $+                "checkExp UpdateSection: cannot project field "+                  <> prettyString f+                  <> " from "+                  <> prettyString t'+        UpdateStepSlice slice -> do+          slice' <- checkSlice slice+          (t', retext) <- sliceShape Nothing slice' =<< normTypeFully t+          (rest', target_t, retext_rest) <- checkSectionSteps t' rest+          pure (UpdateStepSlice slice' : rest', target_t, retext <> retext_rest)+checkExp (AppExp (Loop _ mergepat loopinit form loopbody loc) _) = do+  ((sparams, mergepat', loopinit', form', loopbody'), appres) <-+    checkLoop checkExp (mergepat, loopinit, form, loopbody) loc+  pure $+    AppExp+      (Loop sparams mergepat' loopinit' form' loopbody' loc)+      (Info appres)+checkExp (Constr name es (Info t) loc) = do+  -- The sizes are absorbable: those of the payloads of the other+  -- constructors (and any not determined by the arguments) are+  -- adopted from the context, like the sizes of a hole. See Note+  -- [Size Inference].+  t' <- replaceTyVarsAbsorbable loc t+  es' <- mapM checkExp es+  case t' of+    Scalar (Sum cs)+      | Just name_ts <- M.lookup name cs ->+          zipWithM_ (unify $ mkUsage loc "inferred variant") name_ts $+            map typeOf es'+    _ ->+      error $ "checkExp Constr: " <> prettyString t'+  pure $ Constr name es' (Info t') loc+checkExp (AppExp (If e1 e2 e3 loc) _) = do+  e1' <- checkExp e1+  e2' <- checkExp e2+  e3' <- checkExp e3+  (t, retext) <- unifyBranches loc e2' e3'+  pure $ AppExp (If e1' e2' e3' loc) (Info $ AppRes t retext)+checkExp (AppExp (Match e cs loc) _) = do+  e' <- checkExp e+  mt <- expType e'+  (cs', t, retext) <- checkCases mt cs++  pure $ AppExp (Match e' cs' loc) (Info $ AppRes t retext)+checkExp (Attr info e loc) =+  Attr <$> checkAttr info <*> checkExp e <*> pure loc++checkCase ::+  StructType ->+  CaseBase Info VName ->+  TermTypeM (CaseBase Info VName, StructType, [VName])+checkCase mt (CasePat p e loc) =+  bindingPat [] p mt $ \p' -> do+    e' <- checkExp e+    e_t <- expTypeFully e'+    (e_t', retext) <- unscopeType loc (patNames p') e_t+    pure (CasePat (fmap toStruct p') e' loc, e_t', retext)++updateFieldPath ::+  (Pretty a, Located a) =>+  a ->+  [Name] ->+  StructType ->+  StructType ->+  TermTypeM StructType+updateFieldPath src all_fs ve_t = recurse [] all_fs+  where+    recurse seen [] t = do+      (t', _) <- allDimsFreshInType usage Nonrigid "any" t+      onFailure (CheckingRecordUpdate seen t' ve_t) $+        unify usage t' ve_t+      pure ve_t+      where+        usage = mkUsage (locOf src) "record update"+    recurse seen (f : fs) (Scalar (Record m))+      | Just f_t <- M.lookup f m = do+          f_t' <- recurse (seen ++ [f]) fs f_t+          pure $ Scalar $ Record $ M.insert f f_t' m+    recurse _ _ _ =+      typeError (locOf src) mempty . withIndexLink "record-type-not-known" $+        "Full type of"+          </> indent 2 (pretty src)+          </> textwrap " is not known at this point.  Add a type annotation to the original record to disambiguate."++checkUpdateSteps ::+  SrcLoc ->+  StructType ->+  [UpdateStep Info VName] ->+  TermTypeM ([UpdateStep Info VName], StructType)+checkUpdateSteps _ t [] =+  pure ([], t)+checkUpdateSteps loc t (step : rest) =+  case step of+    UpdateStepSlice slice -> do+      slice' <- checkSlice slice+      (elem_t, _) <- sliceShape (Just (loc, Nonrigid)) slice' =<< normTypeFully t+      (rest', target_t) <- checkUpdateSteps loc elem_t rest+      pure (UpdateStepSlice slice' : rest', target_t)+    UpdateStepField f -> do+      t' <- normTypeFully t+      case t' of+        Scalar (Record fs) | Just f_t <- M.lookup f fs -> do+          (rest', target_t) <- checkUpdateSteps loc f_t rest+          pure (UpdateStepField f : rest', target_t)+        _ -> error $ "checkUpdateSteps: " <> show t'++checkCases ::+  StructType ->+  NE.NonEmpty (CaseBase Info VName) ->+  TermTypeM (NE.NonEmpty (CaseBase Info VName), StructType, [VName])+checkCases mt rest_cs =+  case NE.uncons rest_cs of+    (c, Nothing) -> do+      (c', t, retext) <- checkCase mt c+      pure (NE.singleton c', t, retext)+    (c, Just cs) -> do+      ((c', c_t, _), (cs', cs_t, _)) <-+        (,) <$> checkCase mt c <*> checkCases mt cs+      (brancht, retext) <- unifyBranchTypes (srclocOf c) c_t cs_t+      pure (NE.cons c' cs', brancht, retext)++checkSlice :: SliceBase Info VName -> TermTypeM [DimIndex]+checkSlice = mapM checkDimIndex+  where+    checkDimIndex (DimFix i) =+      DimFix <$> checkExp i+    checkDimIndex (DimSlice i j s) =+      DimSlice <$> traverse checkExp i <*> traverse checkExp j <*> traverse checkExp s++instantiateDimsInReturnType ::+  SrcLoc ->+  Maybe (QualName VName) ->+  ResRetType ->+  TermTypeM (ResType, [VName])+instantiateDimsInReturnType loc fname (RetType dims t)+  | null dims =+      pure (t, mempty)+  | otherwise = do+      dims' <- mapM new dims+      pure (first (onDim $ zip dims $ map (ExpSubst . (`sizeFromName` loc) . qualName) dims') t, dims')+  where+    new =+      newRigidDim loc (RigidRet fname)+        . nameFromText+        . T.takeWhile isAscii+        . baseText+    onDim dims' = applySubst (`lookup` dims')++-- Some information about the function/operator we are trying to+-- apply, and how many arguments it has previously accepted.  Used for+-- generating nicer type errors.+type ApplyOp = (Maybe (QualName VName), Int)++-- | Extract all those names that are bound inside the type.+boundInsideType :: TypeBase Size as -> S.Set VName+boundInsideType (Array _ _ t) = boundInsideType (Scalar t)+boundInsideType (Scalar Prim {}) = mempty+boundInsideType (Scalar (TypeVar _ _ targs)) = foldMap f targs+  where+    f (TypeArgType t) = boundInsideType t+    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))) =+  pn' <> boundInsideType t1 <> S.fromList dims <> boundInsideType t2+  where+    pn' = case pn of+      Unnamed -> mempty+      Named v -> S.singleton v++-- Returns the sizes of the immediate type produced,+-- the sizes of parameter types, and the sizes of return types.+dimUses :: TypeBase Size u -> (Names, Names)+dimUses = flip execState mempty . traverseDims f+  where+    f bound pos e =+      case pos of+        PosImmediate ->+          modify ((fvVars fv, mempty) <>)+        PosParam ->+          modify ((mempty, fvVars fv) <>)+        PosReturn -> pure ()+      where+        fv = freeInExp e `freeWithout` bound++checkApply ::+  SrcLoc ->+  ApplyOp ->+  StructType ->+  Exp ->+  TermTypeM (StructType, StructType, Maybe VName, [VName])+checkApply loc (fname, _) (Scalar (Arrow _ pname _ tp1 tp2)) argexp = do+  let argtype = typeOf argexp+  onFailure (CheckingApply fname argexp tp1 argtype) $ do+    unify (mkUsage argexp "use as function argument") tp1 argtype++    -- Perform substitutions of instantiated variables in the types.+    (tp2_inst, ext) <- instantiateDimsInReturnType loc fname =<< normTypeFully tp2+    argtype' <- normTypeFully argtype++    -- Unification against the argument type may have determined that+    -- some instantiated sizes are existential. Their occurrences in+    -- the return type are replaced with fresh rigid sizes, bound at+    -- the innermost possible position; those bound at the top level+    -- become existentials of the application. The pending size+    -- variables themselves are left alone; occurrences of them+    -- remaining in the AST are existentially bound by+    -- 'bindExistentialInsts' at the end.+    constraints <- getConstraints+    let (inst_pending, inst_reps) = pendingInstSizes constraints+        repOf v = ExpSubst . flip sizeFromName (srclocOf loc) . qualName <$> M.lookup v inst_reps+        tp2_subst = applySubst repOf tp2_inst+    (tp2', inst_ext) <-+      -- 'sizeFree' can only change the type if a pending instantiated size+      -- occurs free in it, so check for a fast path.+      if any inst_pending (fvVars (freeInType tp2_subst))+        then sizeFree loc (find inst_pending . fvVars . freeInExp) tp2_subst+        else pure (tp2_subst, [])+    let ext' = ext <> inst_ext++    -- Check whether this would produce an impossible return type.+    let (tp2_produced_dims, tp2_paramdims) = dimUses tp2'+        problematic = S.fromList ext' <> boundInsideType argtype'+        problem = any (`S.member` problematic) (tp2_paramdims `S.difference` tp2_produced_dims)+    when (not (S.null problematic) && problem) $ do+      typeError loc mempty . withIndexLink "existential-param-ret" $+        "Existential size would appear in function parameter of return type:"+          </> indent 2 (pretty (RetType ext' tp2'))+          </> textwrap "This is usually because a higher-order function is used with functional arguments that return existential sizes or locally named sizes, which are then used as parameters of other function arguments."++    (argext, tp2'') <-+      case pname of+        Named pname'+          | S.member pname' (fvVars $ freeInType tp2') ->+              if hasBinding argexp+                then do+                  warn (srclocOf argexp) $+                    withIndexLink+                      "size-expression-bind"+                      "Size expression with binding is replaced by unknown size."+                  d <- newRigidDim argexp (RigidArg fname $ prettyTextOneLine $ bareExp argexp) "n"+                  let parsubst v =+                        if v == pname'+                          then Just $ ExpSubst $ sizeFromName (qualName d) $ srclocOf argexp+                          else Nothing+                  pure (Just d, applySubst parsubst $ toStruct tp2')+                else+                  let parsubst v =+                        if v == pname'+                          then Just $ ExpSubst $ fromMaybe argexp $ stripExp argexp+                          else Nothing+                   in pure (Nothing, applySubst parsubst $ toStruct tp2')+        _ -> pure (Nothing, toStruct tp2')++    pure (tp1, tp2'', argext, ext')+checkApply _ _ _ _ =+  error "checkApply: array"++-- | Type-check a single expression in isolation.  This expression may+-- turn out to be polymorphic, in which case the list of type+-- parameters will be non-empty.+checkOneExp :: ExpBase NoInfo VName -> TypeM ([TypeParam], Exp)+checkOneExp e = do+  (maybe_tysubsts, e') <- Unsized.checkSingleExp e+  case maybe_tysubsts of+    Left err -> throwError err+    Right (generalised, tysubsts) -> runTermTypeM checkExp tysubsts $ do+      e'' <- checkExp e'+      let t = typeOf e''+      (tparams, _, _) <-+        letGeneralise (nameFromString "<exp>") (srclocOf e) generalised [] $ toRes Nonunique t+      detectAmbiguousSizes+      e''' <- bindExistentialInsts =<< normTypeFully e''+      localChecks tparams e'''+      causalityCheck e'''+      pure (tparams, e''')++-- | Type-check a single size expression in isolation.  This expression may+-- turn out to be polymorphic, in which case it is unified with i64.+checkSizeExp :: ExpBase NoInfo VName -> TypeM Exp+checkSizeExp e = do+  (maybe_tysubsts, e') <- Unsized.checkSizeExp e+  case maybe_tysubsts of+    Left err -> throwError err+    Right (_generalised, tysubsts) -> runTermTypeM checkExp tysubsts $ do+      e'' <- checkExp e'+      when (hasBinding e'') $+        typeError (srclocOf e'') mempty . withIndexLink "size-expression-bind" $+          "Size expression with binding is forbidden."+      normTypeFully e''++-- Verify that all sum type constructors and empty array literals have+-- a size that is known (rigid or a type parameter).  This is to+-- ensure that we can actually determine their shape at run-time.+causalityCheck :: Exp -> TermTypeM ()+causalityCheck binding_body = do+  constraints <- getConstraints++  let checkCausality what known t loc+        | (d, dloc) : _ <-+            mapMaybe (unknown constraints known) $+              S.toList (fvVars $ freeInType t) =+            Just $ lift $ causality what (locOf loc) d dloc t+        | otherwise = Nothing++      checkParamCausality known p =+        checkCausality (pretty p) known (patternType p) (locOf p)++      collectingNewKnown = lift . flip execStateT mempty++      onExp ::+        S.Set VName ->+        Exp ->+        StateT (S.Set VName) (Either TypeError) Exp++      onExp known (Var v (Info t) loc)+        | Just bad <- checkCausality (dquotes (pretty v)) known t loc =+            bad+      onExp known (UpdateSection _ (Info t) loc)+        | Just bad <- checkCausality "projection section" known t loc =+            bad+      onExp known (OpSectionRight _ (Info t) _ _ _ loc)+        | Just bad <- checkCausality "operator section" known t loc =+            bad+      onExp known (OpSectionLeft _ (Info t) _ _ _ loc)+        | Just bad <- checkCausality "operator section" known t loc =+            bad+      onExp known (ArrayLit [] (Info t) loc)+        | Just bad <- checkCausality "empty array" known t loc =+            bad+      onExp known (Hole (Info t) loc)+        | Just bad <- checkCausality "hole" known t loc =+            bad+      onExp known e@(Lambda params body _ _ _)+        | bad : _ <- mapMaybe (checkParamCausality known) params =+            bad+        | otherwise = do+            -- Existentials coming into existence in the lambda body+            -- are not known outside of it.+            void $ collectingNewKnown $ onExp known body+            pure e+      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 (Match scrutinee cs _) (Info res)) = do+        new_known <- collectingNewKnown $ onExp known scrutinee+        void $ recurse (new_known <> known) cs+        modify ((new_known <> S.fromList (appResExt res)) <>)+        pure e+      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 <- collectingNewKnown $ onExp known' x+            void $ seqArgs (new_known <> known') xs+            modify ((new_known <> S.fromList (maybeToList p)) <>)+      onExp known e@(Constr v args (Info t) loc) = do+        seqArgs known args+        pure e+        where+          seqArgs known' []+            | Just bad <- checkCausality (dquotes ("#" <> pretty v)) known' t loc =+                bad+            | otherwise =+                pure ()+          seqArgs known' (x : xs) = do+            new_known <- collectingNewKnown $ onExp known' x+            void $ seqArgs (new_known <> known') xs+            modify (new_known <>)+      onExp+        known+        e@(AppExp (BinOp (f, floc) ft (x, Info xp) (y, Info yp) _) (Info res)) = do+          args_known <-+            collectingNewKnown $ sequencePoint known x y $ catMaybes [xp, yp]+          void $ onExp (args_known <> known) (Var f ft floc)+          modify ((args_known <> S.fromList (appResExt res)) <>)+          pure e+      onExp known e@(AppExp e' (Info res)) = do+        recurse known e'+        modify (<> S.fromList (appResExt res))+        pure e+      onExp known e = do+        recurse known e+        pure e++      recurse known = void . astMap mapper+        where+          mapper = identityMapper {mapOnExp = onExp known}++      sequencePoint known x y ext = do+        new_known <- collectingNewKnown $ onExp known x+        void $ onExp (new_known <> known) y+        modify ((new_known <> S.fromList ext) <>)++  either throwError (const $ pure ()) $+    evalStateT (onExp mempty binding_body) mempty+  where+    unknown constraints known v = do+      guard $ v `S.notMember` known+      loc <- case snd <$> M.lookup v constraints of+        Just (UnknownSize loc _) -> Just loc+        _ -> Nothing+      pure (v, loc)++    causality what loc d dloc t =+      Left . TypeError loc mempty . withIndexLink "causality-check" $+        "Causality check: size"+          <+> dquotes (prettyName d)+          <+> "needed for type of"+          <+> what+          <> colon+            </> indent 2 (pretty t)+            </> "But"+            <+> dquotes (prettyName d)+            <+> "is computed at"+            <+> pretty (locStrRel loc dloc)+          <> "."+            </> ""+            </> "Hint:"+            <+> align+              ( textwrap "Bind the expression producing"+                  <+> dquotes (prettyName d)+                  <+> "with 'let' beforehand."+              )++mustBeIrrefutable :: (MonadTypeChecker f) => Pat StructType -> f ()+mustBeIrrefutable p = do+  case unmatched [p] of+    [] -> pure ()+    ps' ->+      typeError p mempty . withIndexLink "refutable-pattern" $+        "Refutable pattern not allowed here.\nUnmatched cases:"+          </> indent 2 (stack (map pretty ps'))++supportsEquality :: TypeBase dim u -> Bool+supportsEquality (Array _ _ t) = supportsEquality $ Scalar t+supportsEquality (Scalar Prim {}) = True+supportsEquality (Scalar TypeVar {}) = False+supportsEquality (Scalar (Record fs)) = all supportsEquality fs+supportsEquality (Scalar (Sum fs)) = all (all supportsEquality) fs+supportsEquality (Scalar Arrow {}) = False++-- | Check that a type is non-functional, looking up the liftedness of type+-- variables.+orderZeroM :: [TypeParam] -> StructType -> TermTypeM Bool+orderZeroM tparams t = do+  (orderZero t &&) . and <$> mapM isUnlifted (typeQualVars t)+  where+    isUnlifted qv = do+      case find ((== qualLeaf qv) . typeParamName) tparams of+        Just (TypeParamType l _ _) -> pure $ l < Lifted+        _ -> (< Lifted) <$> lookupAbsTy qv++-- | Traverse the expression, emitting warnings and errors for various+-- problems:+--+-- * Unmatched cases.+--+-- * If any of the literals overflow their inferred types. Note:+--  currently unable to detect float underflow (such as 1e-400 -> 0)+--+-- * Function types appearing in places where they are not allowed (e.g.+--   returned from branches), and more generally lifted types used as+--   array elements.+--+-- The rationale is that it is easier to check for these things after all of the+-- type inference has been done, as they complicate the logic. Further, it is+-- also easier to produce good error messages here. The key is that we can only+-- enforce rules that do not affect type inference.+localChecks :: [TypeParam] -> Exp -> TermTypeM ()+localChecks tparams orig_body = void $ check orig_body+  where+    check e@(AppExp (Match _ cs loc) (Info rt)) = do+      ok <- orderZeroM tparams (appResType rt)+      unless ok . typeError loc mempty $+        "Match-expression returns type"+          </> indent 2 (align (pretty (appResType rt)))+          </> "but match-results may not be of function type."+      let ps = fmap (\(CasePat p _ _) -> p) cs+      case unmatched $ NE.toList ps of+        [] -> recurse e+        ps' ->+          typeError loc mempty . withIndexLink "unmatched-cases" $+            "Unmatched cases in match expression:"+              </> indent 2 (stack (map pretty ps'))+    check e@(AppExp (If _ _ _ loc) (Info rt)) = do+      ok <- orderZeroM tparams (appResType rt)+      unless ok . typeError loc mempty $+        "If-expression returns type"+          </> indent 2 (align (pretty (appResType rt)))+          </> "but if-results may not be of function type."+      recurse e+    check e@(ArrayLit _ (Info t) loc) = do+      mapM_ (checkArrayElem loc) $ peelArray 1 t+      recurse e+    check e@(AppExp (LetPat _ p _ _ _) _) =+      mustBeIrrefutable p *> recurse e+    check e@(AppExp (BinOp (v, loc) _ (x, _) _ _) _)+      | qualLeaf v == intrinsicVar "==" = do+          case typeOf x of+            Array {} -> do+              warn loc $+                textwrap+                  "Comparing arrays with \"==\" is deprecated and will stop working in a future revision of the language."+            _ -> pure ()+          checkEquality loc (typeOf x) *> recurse e+    check e@(Var v (Info t) loc)+      | qualLeaf v == intrinsicVar "==" = do+          checkEquality loc t *> recurse e+    check e@(Lambda ps _ _ _ _) =+      mapM_ (mustBeIrrefutable . fmap toStruct) ps *> recurse e+    check e@(AppExp (LetFun _ (tparams', ps, _, _, e1) e2 _) _) = do+      mapM_ (mustBeIrrefutable . fmap toStruct) ps+      localChecks (tparams' <> tparams) e1+      void $ check e2+      pure e+    check e@(AppExp (Loop _ p _ form _ _) _) = do+      mustBeIrrefutable (fmap toStruct p)+      case form of+        ForIn form_p _ -> mustBeIrrefutable form_p+        _ -> pure ()+      ok <- orderZeroM tparams (patternStructType p)+      unless ok . typeError (locOf p) mempty $+        "Loop parameter inferred to have type"+          </> indent 2 (align (pretty p))+          </> "but a loop parameter may not be of function type."+      recurse e+    check e@(IntLit x ty loc) =+      e <$ case ty of+        Info (Scalar (Prim t)) -> errorBounds (inBoundsI x t) x t loc+        _ -> error "Inferred type of int literal is not a number"+    check e@(FloatLit x ty loc) =+      e <$ case ty of+        Info (Scalar (Prim (FloatType t))) -> errorBounds (inBoundsF x t) x t loc+        _ -> error "Inferred type of float literal is not a float"+    check e@(Negate (IntLit x ty loc1) loc2) =+      e <$ case ty of+        Info (Scalar (Prim t)) -> errorBounds (inBoundsI (-x) t) (-x) t (loc1 <> loc2)+        _ -> error "Inferred type of int literal is not a number"+    check e = recurse e+    recurse = astMap identityMapper {mapOnExp = check}++    checkEquality loc t =+      unless (supportsEquality t) $+        typeError loc mempty $+          "Comparing equality of values of type"+            </> indent 2 (pretty t)+            </> "which does not support equality."++    -- Array elements must be unlifted: of non-varying size, and in+    -- particular not functions. This is a stricter requirement than+    -- 'orderZeroM', which permits size-lifted type parameters.+    checkArrayElem loc et = do+      unless (orderZero et) . typeError loc mempty $+        "Type" </> indent 2 (pretty et) </> "found to be functional."+      mapM_ checkElemVar $ typeQualVars et+      where+        checkElemVar qv = do+          l <- case find ((== qualLeaf qv) . typeParamName) tparams of+            Just (TypeParamType l _ tploc) ->+              pure $ Left (l, locOf tploc)+            _ -> Right <$> lookupAbsTy qv+          case l of+            Left (l', tploc)+              | l' /= Unlifted ->+                  typeError loc mempty $+                    "Type parameter"+                      <+> dquotes (pretty qv)+                      <+> "bound at"+                      <+> pretty (locStr tploc)+                      <+> "is lifted and cannot be an array element."+            Right l'+              | l' /= Unlifted ->+                  typeError loc mempty $+                    "Type"+                      <+> dquotes (pretty qv)+                      <+> "is lifted and cannot be an array element."+            _ -> pure ()++    bitWidth ty = 8 * intByteSize ty :: Int++    inBoundsI x (Signed t) = x >= -2 ^ (bitWidth t - 1) && x < 2 ^ (bitWidth t - 1)+    inBoundsI x (Unsigned t) = x >= 0 && x < 2 ^ bitWidth t+    inBoundsI x (FloatType Float16) = not $ isInfinite (fromIntegral x :: Half)+    inBoundsI x (FloatType Float32) = not $ isInfinite (fromIntegral x :: Float)+    inBoundsI x (FloatType Float64) = not $ isInfinite (fromIntegral x :: Double)+    inBoundsI _ Bool = error "Inferred type of int literal is not a number"+    inBoundsF x Float16 = not $ isInfinite (realToFrac x :: Float)+    inBoundsF x Float32 = not $ isInfinite (realToFrac x :: Float)+    inBoundsF x Float64 = not $ isInfinite x++    errorBounds inBounds x ty loc =+      unless inBounds $+        typeError loc mempty . withIndexLink "literal-out-of-bounds" $+          "Literal "+            <> pretty x+            <> " out of bounds for inferred type "+            <> pretty ty+            <> "."++-- | Check restrictions on recursive functions: the result must be first-order,+-- and any recursive applications must be invariant with respect to the+-- higher-order arguments. These are syntactic checks.+--+-- 'fname', 'params' and 'ret' describe the function whose body this is; if the+-- function is not recursive, this is a no-op.+recursionCheck ::+  [TypeParam] -> VName -> [Pat ParamType] -> ResType -> SrcLoc -> Exp -> TermTypeM ()+recursionCheck tparams fname params ret fun_loc body =+  when (fname `S.member` fvVars (freeInExp body)) $ do+    checkRet+    higher_order <-+      mapM (fmap not . orderZeroM tparams . patternStructType) params+    when (or higher_order) $ void $ check higher_order body+  where+    checkRet = do+      ok <- orderZeroM tparams $ toStruct ret+      unless ok . typeError fun_loc mempty $+        "Recursive function"+          <+> dquotes (prettyName fname)+          <+> "returns type"+          </> indent 2 (pretty ret)+          </> "which is not first-order."+          </> "Write the function type as further parameters instead."++    check ho e@(AppExp (Apply f args _) _)+      | Var v _ _ <- f,+        qualLeaf v == fname = do+          checkRecApply ho (locOf e) $ map snd $ NE.toList args+          mapM_ (check ho . snd) args+          pure e+    check _ (Var v _ loc)+      | qualLeaf v == fname =+          typeError loc mempty $+            "Recursive reference to"+              <+> dquotes (prettyName fname)+              <+> "must be a fully saturated application, as it has a higher-order parameter."+    check ho e = recurse ho e+    recurse ho = astMap identityMapper {mapOnExp = check ho}++    checkRecApply ho loc args = do+      unless (length args == length params) . typeError loc mempty $+        "Recursive application of"+          <+> dquotes (prettyName fname)+          <+> "is not fully saturated: expected"+          <+> pretty (length params)+          <+> "arguments but got"+          <+> pretty (length args)+          <> "."+      forM_ (zip3 ho params args) $ \(is_ho, p, arg) ->+        when is_ho $+          case arg of+            Var v _ _ | qualLeaf v `elem` patNames p -> pure ()+            _ ->+              typeError (locOf arg) mempty $+                "Higher-order argument in recursive application of"+                  <+> dquotes (prettyName fname)+                  <+> "must be passed unchanged, i.e. be the corresponding parameter."++-- | Instantiated sizes that unification has determined to be existential+-- ("pending"), and a mapping from pending copies to a representative: copies+-- from the same occurrence of an instantiated type parameter that were absorbed+-- from the same source denote the same existential size. See Note [Size+-- Inference].+pendingInstSizes :: Constraints -> (VName -> Bool, M.Map VName VName)+pendingInstSizes constraints = (pending, reps)+  where+    key v = case snd <$> M.lookup v constraints of+      Just (ExistentialSize k _ _) -> Just k+      _ -> Nothing+    pending v = case snd <$> M.lookup v constraints of+      Just ExistentialSize {} -> True+      Just (CopySize c _ _) -> isJust $ key c+      _ -> False+    groups =+      M.fromListWith+        (<>)+        [ ((occ, k), [v])+        | (v, (_, CopySize c occ _)) <- M.toList constraints,+          Just (Just k) <- [key c]+        ]+    reps =+      M.fromList+        [ (v, rep)+        | vs <- M.elems groups,+          let rep = minimum vs,+          v <- vs,+          v /= rep+        ]++-- | Instantiated sizes (at or above the given level) that are still pending and+-- have not been determined to be existential behave like ordinary sizes from+-- here on: canonical sizes are plain size variables, and copies are equal to+-- their canonical variable. Existential ones are left alone; they are handled+-- by 'bindExistentialInsts'.+collapseInstSizes :: Level -> TermTypeM ()+collapseInstSizes min_lvl = do+  constraints <- getConstraints+  let nonExistential c = case snd <$> M.lookup c constraints of+        Just ExistentialSize {} -> False+        _ -> True+      collapse (lvl, CopySize c _ usage)+        | lvl >= min_lvl,+          nonExistential c =+            (lvl, Size (Just $ sizeFromName (qualName c) $ srclocOf usage) usage)+      collapse (lvl, InstSize _ usage)+        | lvl >= min_lvl = (lvl, Size Nothing usage)+      collapse x = x+  modifyConstraints $ M.map collapse++-- | Instantiated sizes that unification determined to be existential may remain+-- free in some types recorded in the AST - in particular the instantiated types+-- of higher-order functions, where the existential size occurs in the return+-- type of a function-typed parameter. Existentially bind such sizes at the+-- innermost possible position, mirroring what the type of the function argument+-- looks like.+bindExistentialInsts :: (ASTMappable e) => e -> TermTypeM e+bindExistentialInsts x = do+  -- 'pendingInstSizes' scans the entire constraint set, but its result+  -- is invariant across this traversal (any rigid sizes we introduce+  -- below are not instantiated sizes), so we compute it once instead of+  -- once per type in the AST.+  constraints <- getConstraints+  let (pending, reps) = pendingInstSizes constraints+      repOf v = ExpSubst . flip sizeFromName mempty . qualName <$> M.lookup v reps+      relevant v = pending v || v `M.member` reps++      onType ::+        (Substitutable (TypeBase Size u)) =>+        TypeBase Size u ->+        TermTypeM (TypeBase Size u, [VName])+      onType t+        -- Fast path: this type mentions no pending or copied+        -- instantiated size, so 'applySubst'/'sizeFree' would be+        -- no-ops. Most types take this path.+        | not (any relevant $ fvVars $ freeInType t) = pure (t, [])+        | otherwise =+            sizeFree mempty (find pending . fvVars . freeInExp) $ applySubst repOf t++      onStruct ::+        (Substitutable (TypeBase Size u)) =>+        TypeBase Size u ->+        TermTypeM (TypeBase Size u)+      onStruct t = do+        (t', ext) <- onType t+        -- Existential sizes at the top level of a type have nowhere to+        -- be bound. Those that absorbed a rigid unknown size stand for+        -- a size that is actually computed at the recorded location,+        -- so they become rigid unknown sizes there, subjecting them to+        -- the causality check. The rest (absorbed from declared+        -- existentials, e.g. by a hole) are left alone.+        if null ext+          then pure t'+          else do+            let computedAt v = case snd <$> M.lookup v constraints of+                  Just (ExistentialSize _ mloc _) -> mloc+                  Just (CopySize c _ _)+                    | Just (ExistentialSize _ mloc _) <- snd <$> M.lookup c constraints ->+                        mloc+                  _ -> Nothing+            repls <- fmap (M.fromList . catMaybes) . forM (S.toList $ fvVars $ freeInType t) $ \v ->+              case computedAt v of+                Just dloc -> do+                  v' <- newRigidDim dloc (RigidRet Nothing) "d"+                  pure $ Just (v, ExpSubst $ sizeFromName (qualName v') $ srclocOf dloc)+                Nothing -> pure Nothing+            pure $ applySubst (`M.lookup` repls) t++      tv =+        ASTMapper+          { mapOnExp = astMap tv,+            mapOnName = pure,+            mapOnStructType = onStruct,+            mapOnParamType = onStruct,+            mapOnResRetType = \(RetType dims t) -> do+              (t', ext) <- onType t+              pure $ RetType (dims <> ext) t'+          }+  -- Global fast path: with no instantiated-size constraints at all, 'relevant'+  -- is false everywhere, so the traversal would just rebuild an identical copy.+  -- Skip it entirely - this is the common case.+  if any (isInstSize . snd) constraints+    then astMap tv x+    else pure x+  where+    isInstSize ExistentialSize {} = True+    isInstSize CopySize {} = True+    isInstSize _ = False++detectAmbiguousSizes :: TermTypeM ()+detectAmbiguousSizes = do+  collapseInstSizes 0+  constraints <- getConstraints+  mapM_ (notice constraints) $ M.toList constraints+  where+    -- Sizes that arise from instantiating inferred types have+    -- uninformative provenance. If a size variable with better+    -- provenance (e.g. a source-level size binder, or an+    -- instantiated size parameter) has been unified with the+    -- ambiguous size, we report that variable instead. This affects+    -- only the error message; which sizes are ambiguous is already+    -- settled.+    uninformative (Usage Nothing _) = True+    uninformative (Usage (Just u) _) = u `elem` ["replaceTyVars", "instantiation"]++    chase constraints w = case snd <$> M.lookup w constraints of+      Just (Size (Just (Var w' _ _)) _) -> chase constraints (qualLeaf w')+      _ -> w++    blame constraints v usage+      | uninformative usage,+        (v', usage') : _ <-+          [ (w, w_usage)+          | (w, (_, Size (Just (Var w1 _ _)) w_usage)) <- M.toList constraints,+            not $ uninformative w_usage,+            chase constraints (qualLeaf w1) == v+          ] =+          (v', usage')+      | otherwise = (v, usage)++    notice constraints (v, (_, Size Nothing usage)) =+      case blame constraints v usage of+        (v', Usage Nothing loc) ->+          typeError loc mempty . withIndexLink "ambiguous-size" $+            "Ambiguous size" <+> dquotes (prettyName v') <> "."+        (v', Usage (Just u) loc) ->+          typeError loc mempty . withIndexLink "ambiguous-size" $+            "Ambiguous size" <+> dquotes (prettyName v') <+> "arising from" <+> pretty u <> "."+    notice _ _ = pure ()++-- | The names bound by these parameter patterns that are not the name of+-- a parameter itself, and hence cannot occur in the function type.+hiddenParamNames :: [Pat ParamType] -> [VName]+hiddenParamNames params = hidden+  where+    param_all_names = mconcat $ map patNames params+    named (Named x, _, _) = Just x+    named (Unnamed, _, _) = Nothing+    param_names =+      S.fromList $ mapMaybe (named . patternParam) params+    hidden = filter (`notElem` param_names) param_all_names++-- | Rename the sizes bound by a type (parameter names and existential+-- quantifiers) to fresh names.+renameTypeBinders :: (Monoid u) => TypeBase Size u -> TermTypeM (TypeBase Size u)+renameTypeBinders (Scalar (Arrow u pn d pt (RetType dims rt))) = do+  pt' <- renameTypeBinders pt+  (pn', pn_subst) <- case pn of+    Named v -> do+      v' <- newName v+      pure (Named v', M.singleton v v')+    Unnamed -> pure (Unnamed, mempty)+  dims' <- mapM newName dims+  let subst = pn_subst <> M.fromList (zip dims dims')+      toSize v = ExpSubst $ sizeFromName (qualName v) mempty+  rt' <- renameTypeBinders $ applySubst (fmap toSize . (`M.lookup` subst)) rt+  pure $ Scalar $ Arrow u pn' d pt' $ RetType dims' rt'+renameTypeBinders (Scalar (Record fs)) =+  Scalar . Record <$> traverse renameTypeBinders fs+renameTypeBinders (Scalar (Sum cs)) =+  Scalar . Sum <$> traverse (traverse renameTypeBinders) cs+renameTypeBinders t = pure t++checkBinding ::+  ( VName,+    Maybe (TypeExp Exp VName),+    [TypeParam],+    [PatBase Info VName ParamType],+    ExpBase Info VName,+    SrcLoc+  ) ->+  TermTypeM+    ( [TypeParam],+      [Pat ParamType],+      Maybe (TypeExp Exp VName),+      ResRetType,+      Exp+    )+checkBinding (fname, maybe_retdecl, tparams, params, body, loc) =+  incLevel . bindingParams tparams params $ \params' -> do+    maybe_retdecl' <- traverse checkTypeExpNonrigid maybe_retdecl++    -- Bind the name in scope of its own body so it may recurse. Harmless even+    -- when the function is not actually recursive, as name resolution has+    -- hooked things up properly anyway. See Note [Checking recursive+    -- functions].+    self_binding <- case maybe_retdecl' of+      Just (_, ret, ext) ->+        BoundV tparams <$> renameTypeBinders (funType params' (RetType ext ret))+      Nothing -> pure RecursiveV+    (body', body_t) <-+      localScope (\scope -> scope {scopeVtable = M.insert fname self_binding $ scopeVtable scope}) $+        checkFunBody+          params'+          body+          ((\(_, x, _) -> x) <$> maybe_retdecl')+          (maybe loc srclocOf maybe_retdecl)++    params'' <- mapM updateTypes params'++    (maybe_retdecl'', rettype) <- case maybe_retdecl' of+      Just (retdecl', ret, _) -> do+        ret' <- normTypeFully ret+        pure (Just retdecl', ret')+      Nothing ->+        pure (Nothing, toRes Nonunique body_t)++    verifyFunctionParams (Just fname) params''++    (tparams', params''', rettype') <-+      letGeneralise (baseName fname) loc tparams params''+        =<< unscopeUnknown rettype++    when+      ( null params+          && any isSizeParam tparams'+          && not (null (retDims rettype'))+      )+      $ typeError loc mempty+      $ textwrap "A size-polymorphic value binding may not have a type with an existential size."+        </> "Type of this binding is:"+        </> indent 2 (pretty rettype')+        </> "with the following type parameters:"+        </> indent 2 (sep $ map pretty $ filter isSizeParam tparams')++    pure (tparams', params''', maybe_retdecl'', rettype', body')++-- | 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+  where+    onParam :: TypeBase Size als -> S.Set VName+    onParam (Scalar Arrow {}) = mempty+    onParam (Scalar (Record fs)) = mconcat $ map onParam $ M.elems fs+    onParam (Scalar (TypeVar _ _ targs)) = mconcat $ map onTypeArg targs+    onParam t = fvVars $ freeInType t+    onTypeArg (TypeArgDim (Var d _ _)) = S.singleton $ qualLeaf d+    onTypeArg (TypeArgDim _) = mempty+    onTypeArg (TypeArgType t) = onParam t+sizeNamesPos _ = mempty++-- | Verify certain restrictions on function parameters, and bail out+-- on dubious constructions.+--+-- These restrictions apply to all functions (anonymous or otherwise).+-- Top-level functions have further restrictions that are checked+-- during let-generalisation.+--+-- The parameters are assumed to already have their types normalised+-- ('updateTypes'), which both callers do immediately beforehand.+verifyFunctionParams :: Maybe VName -> [Pat ParamType] -> TermTypeM ()+verifyFunctionParams fname params =+  onFailure (CheckingParams (baseName <$> fname)) $+    verifyParams (foldMap patNames params) params+  where+    verifyParams forbidden (p : ps)+      | d : _ <- filter (`elem` forbidden) $ S.toList $ fvVars $ freeInPat p =+          typeError p mempty . withIndexLink "inaccessible-size" $+            "Parameter"+              <+> dquotes (pretty p)+              </> "refers to size"+              <+> dquotes (prettyName d)+              <> comma+                </> textwrap "which will not be accessible to the caller"+              <> comma+                </> textwrap "possibly because it is nested in a tuple or record."+                </> textwrap "Consider ascribing an explicit type that does not reference "+              <> dquotes (prettyName d)+              <> "."+      | otherwise = verifyParams forbidden' ps+      where+        forbidden' =+          case patternParam p of+            (Named v, _, _) -> delete v forbidden+            _ -> forbidden+    verifyParams _ [] = pure ()++-- | Move existentials down to the level where they are actually used+-- (i.e. have their "witnesses").  E.g. changes+--+-- @+-- ?[n].bool -> [n]bool+-- @+--+-- to+--+-- @+-- bool -> ?[n].[n]bool+-- @+injectExt :: [VName] -> TypeBase Size u -> RetTypeBase Size u+injectExt [] ret = RetType [] ret+injectExt ext ret = RetType ext_here $ deeper ret+  where+    (immediate, _) = dimUses ret+    (ext_here, ext_there) = partition (`S.member` immediate) ext+    deeper :: TypeBase Size u -> TypeBase Size u+    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 d1 t1 (RetType t2_ext t2))) =+      Scalar $ Arrow als p d1 t1 $ injectExt (nubOrd (ext_there <> t2_ext)) t2+    deeper (Scalar (TypeVar u tn targs)) =+      Scalar $ TypeVar u tn $ map deeperArg targs+    deeper t@Array {} = t++    deeperArg (TypeArgType t) = TypeArgType $ deeper t+    deeperArg (TypeArgDim d) = TypeArgDim d++-- | Find all size variables in the given type that are covered by the+-- constraints, and produce size parameters that close over them.+--+-- The passed-in list of type parameters is always prepended to the+-- produced list of type parameters.+closeOverSizes ::+  Name ->+  SrcLoc ->+  [TypeParam] ->+  [StructType] ->+  ResType ->+  Constraints ->+  TermTypeM ([TypeParam], ResRetType)+closeOverSizes defname defloc tparams paramts ret substs = do+  (more_tparams, retext) <-+    partitionEithers . catMaybes+      <$> mapM closeOver (M.toList $ M.map snd to_close_over)+  let mkExt v =+        case M.lookup v substs of+          Just (_, UnknownSize {}) -> Just v+          _ -> Nothing++  pure+    ( tparams+        ++ more_tparams,+      injectExt (nubOrd $ retext ++ mapMaybe mkExt (S.toList $ fvVars $ freeInType ret)) ret+    )+  where+    -- Diet does not matter here.+    t = foldFunType (map (toParam Observe) paramts) $ RetType [] ret+    visible = typeVars t <> fvVars (freeInType t)+    to_close_over =+      M.filterWithKey (\k _ -> k `S.member` visible) substs++    (produced_sizes, param_sizes) = dimUses t++    -- Avoid duplicate type parameters.+    closeOver (k, _)+      | k `elem` map typeParamName tparams =+          pure Nothing+    closeOver (k, Size Nothing _) =+      pure $ Just $ Left $ TypeParamDim k mempty+    closeOver (k, UnknownSize _ _)+      | k `S.member` param_sizes,+        k `S.notMember` produced_sizes = do+          notes <- dimNotes defloc $ sizeFromName (qualName k) mempty+          typeError defloc notes . withIndexLink "unknown-param-def" $+            "Unknown size"+              <+> dquotes (prettyName k)+              <+> "in parameter of"+              <+> dquotes (prettyName defname)+              <> ", which is inferred as:"+                </> indent 2 (pretty t)+      | k `S.member` produced_sizes =+          pure $ Just $ Right k+    closeOver (_, _) =+      pure Nothing++letGeneralise ::+  Name ->+  SrcLoc ->+  [TypeParam] ->+  [Pat ParamType] ->+  ResType ->+  TermTypeM ([TypeParam], [Pat ParamType], ResRetType)+letGeneralise defname defloc tparams params restype =+  onFailure (CheckingLetGeneralise defname) $ do+    cur_lvl <- curLevel+    collapseInstSizes $ cur_lvl - length params++    -- Re-normalise the types so that any instantiated sizes+    -- collapsed above are expressed in terms of their canonical+    -- variables, which can then be closed over.+    params' <- mapM updateTypes params+    restype' <- normTypeFully restype++    now_substs <- getConstraints++    -- Candidates for let-generalisation are those size variables that+    --+    -- (1) were not known before we checked this function, and+    --+    -- (2) are not used in the (new) definition of any size variables+    -- known before we checked this function.++    -- Criteria (1) and (2) is implemented by looking at the binding+    -- level of the size variables.+    let candidate (lvl, _) = lvl >= (cur_lvl - length params)+        new_substs = M.filter candidate now_substs++    (tparams', RetType ret_dims restype'') <-+      closeOverSizes+        defname+        defloc+        tparams+        (map patternStructType params')+        restype'+        new_substs++    restype''' <- updateTypes restype''++    let used_sizes =+          freeInType restype''' <> foldMap (freeInType . patternType) params'+    case filter ((`S.notMember` fvVars used_sizes) . typeParamName) $+      filter isSizeParam tparams' of+      [] -> pure ()+      tp : _ -> unusedSize $ SizeBinder (typeParamName tp) (srclocOf tp)++    -- We keep those type variables that were not closed over by+    -- let-generalisation.+    modifyConstraints $ M.filterWithKey $ \k _ -> k `notElem` map typeParamName tparams'++    pure (tparams', params', RetType ret_dims restype''')++-- | Check the body of a function, and return it along with its type as+-- seen from outside the function: any 'hiddenParamNames' occurring as+-- sizes are replaced with fresh unknowns, which whoever decides the+-- function's return type then binds existentially.+checkFunBody ::+  [Pat ParamType] ->+  Exp ->+  Maybe ResType ->+  SrcLoc ->+  TermTypeM (Exp, StructType)+checkFunBody params body maybe_rettype loc = do+  body' <- checkExp body+  let hidden = hiddenParamNames params+  (body_t, _) <- unscopeType loc hidden =<< expTypeFully body'++  -- Unify body return type with return annotation, if one exists.+  case maybe_rettype of+    Just rettype ->+      case find (`elem` hidden) $ fvVars $ freeInType rettype of+        Just v ->+          typeError loc mempty $+            "The return type annotation"+              </> indent 2 (align (pretty rettype))+              </> "refers to the name"+              <+> dquotes (prettyName v)+              <+> "which is bound to an inner component of a function parameter."+        Nothing -> do+          let usage = mkUsage body "return type annotation"+          onFailure (CheckingReturn rettype body_t) $+            unify usage (toStruct rettype) body_t+    Nothing -> pure ()++  pure (body', body_t)++-- | Type-check a top-level (or module-level) function definition.+-- Despite the name, this is also used for checking constant+-- definitions, by treating them as 0-ary functions.+checkFunDef ::+  ( VName,+    Maybe (TypeExp (ExpBase NoInfo VName) VName),+    [TypeParam],+    [PatBase NoInfo VName ParamType],+    ExpBase NoInfo VName,+    SrcLoc+  ) ->+  TypeM+    ( [TypeParam],+      [Pat ParamType],+      Maybe (TypeExp Exp VName),+      ResRetType,+      Exp+    )+checkFunDef (fname, retdecl, tparams, params, body, loc) =+  doChecks =<< Unsized.checkValDef (fname, retdecl, tparams, params, body, loc)+  where+    doChecks (maybe_tysubsts, params', retdecl', body') =+      case maybe_tysubsts of+        Left err -> throwError err+        Right (generalised, tysubsts) ->+          runTermTypeM checkExp tysubsts $ do+            (tparams', params'', retdecl'', RetType dims rettype', body'') <-+              checkBinding (fname, retdecl', generalised <> tparams, params', body', loc)++            -- Since this is a top-level function, we also resolve overloaded+            -- types, using either defaults or complaining about ambiguities.+            detectAmbiguousSizes++            -- Then replace all inferred types in the body and parameters.+            body''' <- bindExistentialInsts =<< normTypeFully body''+            params''' <- mapM normTypeFully params''+            retdecl''' <- traverse updateTypes retdecl''+            rettype'' <- normTypeFully rettype'++            -- Check if the function body can actually be evaluated.+            causalityCheck body'''++            -- Check for various problems.+            mapM_ (mustBeIrrefutable . fmap toStruct) params''+            localChecks tparams' body'''+            recursionCheck tparams' fname params''' rettype'' loc body'''++            let ((body'''', updated_ret), errors) =+                  Consumption.checkValDef+                    ( fname,+                      params''',+                      body''',+                      RetType dims rettype'',+                      retdecl''',+                      loc+                    )++            mapM_ throwError errors++            pure (tparams', params''', retdecl''', updated_ret, body'''')++-- Note [Size Inference]+--+-- Type checking of terms is split into two passes. The unsized type checker+-- (Language.Futhark.TypeChecker.Terms.Unsized) infers types while ignoring+-- sizes entirely - its solution maps type variables to types whose dimensions+-- are all vacuous. The sized type checker (this module) receives that solution+-- (the 'termTyVars' field) and is responsible only for inferring sizes: the+-- concrete size of every dimension, and where existential quantifiers go. It+-- never re-infers anything besides sizes and where existential quantifiers go.+--+-- Whenever the size checker needs the type of something that the unsized+-- checker inferred, it instantiates the unsized type by replacing every+-- dimension with a fresh size variable ('instTyVars' when instantiating a type+-- scheme, 'replaceTyVars' elsewhere). Ordinary size unification then determines+-- what these variables stand for. This works out simply enough, except for+-- existential sizes.+--+-- ## Existential sizes+--+-- Consider+--+--   def (|>) '^a '^b (x: a) (f: a -> b) : b = f x+--+--   def main (xs: []i32) = xs |> filter (> 0)+--+-- where "b" is instantiated with the type of "filter (> 0)", which is [n]i32 ->+-- ?[m].[m]i32. At instantiation time we only know the unsized type []i32 for+-- "b" - existential sizes are invisible to the unsized pass, so we cannot know+-- that the size of "b" might be existential to the function. Worse, if the+-- instantiation does turn out to contain an existential size, then every+-- occurrence of "b" in the instantiated type scheme denotes a *distinct*+-- existential:+--+--   [n]i32 -> ([n]i32 -> ?[m].[m]i32) -> ?[m'].[m']i32+--+-- But if the instantiation turns out to have an ordinary size (say "xs |> map+-- (+1)"), all occurrences denote the *same* size, and we must not lose that+-- connection, or we would infer needlessly existential types.+--+-- We address this via three constraint forms (see 'Constraint'):+--+-- - InstSize: a canonical instantiated size, i.e. a dimension of the first+--   occurrence of an instantiated type parameter (or an instantiated size+--   parameter, which has the same nature).+--+-- - CopySize: a dimension of a later occurrence. These are given distinct names+--   precisely so that occurrences can become distinct existentials, but as long+--   as the size is not existential, a copy is equal to its canonical variable+--   (and unification treats it so, by redirecting links to the canonical+--   variable).+--+-- - ExistentialSize: When unification would otherwise fail by linking an+--   instantiated size to a size bound locally in the other type (an existential+--   or a parameter of a function type in the argument), it instead marks the+--   canonical variable with this ('unifySizes'). We say that the instantiated+--   size *absorbs* the locally bound size, and we call size variables that are+--   permitted to do so *absorbable* (see below). Absorption is refused for+--   unlifted type parameters, which cannot have existential sizes.+--+-- The pending existentials are then turned into proper sizes at the places that+-- can bind them:+--+-- - checkApply binds pending sizes in the return type of an application using+--   sizeFree: at the innermost RetType where possible, with the remainder+--   becoming existentials of the application itself (AppRes).+--+-- - bindExistentialInsts does the same for pending sizes that remain in types+--   recorded in the AST, in particular instantiated higher-order function+--   types, where the existential occurs in the return type of a function-typed+--   *parameter* and hence never passes through checkApply's return type.+--+-- - letGeneralise demotes instantiated sizes that are still pending and never+--   became existential to ordinary sizes (collapseInstSizes), at which point+--   they can be closed over as hidden size parameters. For local functions this+--   recreates the per-use size freshness that the old type checker obtained+--   from let-generalising type variables.+--+-- Two occurrences of an instantiated type parameter absorbed from the same+-- source must moreover denote the *same* existential, or we would infer+-- unwitnessed existentials where witnessed ones are possible. This is why+-- ExistentialSize records the size it was unified with and CopySize records+-- which occurrence it belongs to: pending sizes from the same occurrence with+-- the same source are given a single fresh name ('pendingInstSizes' in Terms).+--+-- ## Sizes bound by parameter patterns+--+-- An existential also arises whenever the type of a function body mentions a+-- name that a parameter *pattern* binds without it being the name of the+-- parameter itself, as the "k" of+--+--   \(k: i64, x: i32) -> replicate k x+--+-- where the parameter as a whole is anonymous. Such a name is not bound in the+-- function type, so we replace it with a fresh unknown size ('hiddenParamNames'+-- and 'unscopeType') in the function type. The logic is that we want to infer+-- the type that is visible from the "outside".+--+-- ## Dependent function types+--+-- A related problem is a dependent function type such as (n: i64) -> [n]i32 ->+-- [n]i32, whether reached by instantiating a type parameter with it or by+-- projecting it out of a value (e.g. a record field). The unsized pass preserves+-- parameter names, but the connection between the sizes and the binder is+-- exactly what was erased. Since every fresh size variable occurs exactly once,+-- and binders are cloned between occurrences of the type, it is safe to link a+-- fresh size to a binder of the type itself - such binders are registered+-- whenever the type is reconstructed from the erased solution ('registerBinders',+-- called from both 'instTyVars' and 'replaceTyVars'), and 'unifySizes' permits+-- exactly those links. Linking to any *other* locally bound size is what+-- signifies an existential (see above), or an error for sizes with no such+-- privileges. This binder linking is independent of absorption: it is available+-- even to the non-absorbable ('Unlifted') sizes of 'replaceTyVars', which is why+-- projecting a dependent function out of a record preserves its dependency+-- (tests/shapes/funshape11.fut).+--+-- ## Absorption privileges+--+-- To make the term of art explicit: a fresh size variable is *absorbable* if+-- unification may determine that it stands for an existential size. When an+-- absorbable size meets a size that is bound locally within the type it is+-- unified with, the mismatch is not an error; instead the absorbable size+-- absorbs the locally bound size - it is marked as a pending existential+-- (ExistentialSize), and is eventually existentially bound by the machinery+-- above. A size variable that is not absorbable must be resolved to an ordinary+-- size that is in scope, and encountering a locally bound size is an error for+-- it. Mechanically, absorbable sizes are exactly those constrained by InstSize+-- (or CopySize referring to one).+--+-- Which fresh sizes are absorbable is a fine line:+--+-- - Sizes arising from instantiation - of type parameters and of size+--   parameters - are absorbable. They occur exactly once, and stand for+--   "whatever size the context provides", which may well be existential.+--+-- - The sizes of a hole ('replaceTyVarsAbsorbable') are absorbable, as a hole+--   adopts whatever type the context provides. This includes adopting declared+--   existentials: a flexible existential that unification resolves to an+--   absorbable size is absorbed by it (see the end of the arrow case of+--   'unifyWith').+--+-- - The sizes of the inferred return type of a lambda and of a constructor+--   expression are likewise absorbable: a lambda returns whatever type the+--   context requires (in particular a lambda body ending in "#None" may well+--   have an existential option type), and the payload sizes of a constructor+--   application that are not determined by its arguments - notably those of+--   the *other* constructors - are adopted from the context. Relatedly, the+--   inferred return type of a lambda only hides (existentially binds) sizes+--   that are rigid; a still-flexible size has not been determined yet, and+--   hiding it would sever its connection to whatever the enclosing context+--   determines it to be (see 'inferReturnSizes').+--+-- - Sizes in the types of lambda parameters and patterns ('replaceTyVars') are+--   not absorbable (they are constrained 'InstSize' 'Unlifted', which can link+--   to binders of the type itself but not absorb an existential). They name the+--   sizes of actual bound values, and must be resolved to real sizes. For+--   example, this is what rejects+--+--     def f : (k: i64) -> [k]i32 -> i64 = \_ xs -> length xs+--+--   where the size of "xs" would otherwise silently absorb "k", which is not in+--   scope in the function body (tests/shapes/paramsize1.fut). Contrast with+--   tests/issue1168.fut, where the same shape of program must be accepted+--   because the inner function is size-generalised, and the *instantiation* of+--   its hidden size parameter is what absorbs the bound size of the expected+--   type.+--+-- ## Absorption and causality+--+-- Absorbing an existential is not always innocent. Whether the size that was+-- absorbed is a *declared* existential (of an ascribed type) or a *rigid+-- unknown* size (one whose value is computed at a specific point in the+-- program, such as the result of applying a function with an existential+-- return type) makes a difference:+--+--   def ite b t f = if b then t () else f ()+--   def f : () -> option ([]i32) = \() -> #None                    -- fine+--   def g b = ite b (\() -> #None) (\() -> #Some (gen ()))         -- rejected+--+-- In g, the type of "#None" is forced (via the instantiation of the type+-- parameter of "ite") to have the size produced by "gen ()" inside the other+-- lambda - a size that is only computed elsewhere, so the constructed value+-- cannot know its payload size (tests/sumtypes/sumtype52.fut). ExistentialSize+-- therefore records the location when the absorbed size was rigid, and when+-- such a pending existential remains in a type with no position to bind it,+-- 'bindExistentialInsts' turns it into a rigid unknown size at the recorded+-- location. The causality check then rejects expressions that need it (a sum+-- constructor, say) before it is computed, with no knowledge of this+-- machinery. Two subtleties: the ext variables temporarily introduced in the+-- arrow case of 'unifyWith' shadow the registered constraints of the binders,+-- so the rigidity of a binder is determined from the constraints as they were+-- before ('rigidPre'); and a pending size that checkApply has already bound in+-- a remaining parameter type (of a partially applied function) is registered+-- as a rigid unknown size by 'sizeFree', which is what carries the obligation+-- across applications.++-- Note [Checking recursive functions]+--+-- A function may refer to itself in its own body. The difficulty is that the+-- function's type is not fully known until we have checked that body, so we+-- cannot simply look the name up like any other. We follow the textbook+-- Hindley-Milner treatment on monomoprhic recursion, where all occurrences of+-- the function within its own body share a single type. However, we add the+-- twist that *size parameters* are allowed to differ.+--+-- The handling is spread across name resolution and the two type-checking+-- passes:+--+-- 1. Name resolution brings the function name into scope of its own body, but+--    only for *syntactic* functions (those with parameters).+--+-- 2. The unsized pass binds the name to a fresh monomorphic type variable while+--    checking the body, then emits a single constraint equating that variable+--    with the actual function type. This is standard monomorphic recursion.+--+-- 3. The sized pass considers two cases, distinguished by whether the function+--    has a declared return type:+--+--    * With a declared return type, the function's size-precise type is known+--      before we check the body, so we bind the name to that type scheme+--      ('BoundV'). Recursive occurrences are then instantiated like calls to+--      any other function: sizes are refreshed per occurrence, but the size+--      *relationships* of the signature are kept - e.g. that '[n]i32 -> [n]i32'+--      returns an array the size of its argument. This permits size-polymorphic+--      recursion.+--+--    * Without a declared return type, we cannot know the return type in+--      advance, so we bind the name to 'RecursiveV' . 'lookupVar' then resolves+--      each occurrence with 'replaceTyVars' on the type the unsized pass+--      recorded, which creates *fresh, unrelated* sizes for every dimension,+--      which is OK whenever no size relationship in the return type matters.+--      This means we cannot *infer* size constraints for recursive functions.+--+-- Only top-level self-recursion is handled. Mutual recursion and local+-- (let-bound) recursion are not.
src/Language/Futhark/TypeChecker/Terms/Loop.hs view
@@ -1,6 +1,5 @@--- | Type inference of @loop@.  This is complicated because of the--- uniqueness and size inference, so the implementation is separate--- from the main type checker.+-- | Size inference of @loop@. This is complicated, so the implementation is+-- separate from the main type checker. module Language.Futhark.TypeChecker.Terms.Loop   ( UncheckedLoop,     CheckedLoop,@@ -84,6 +83,20 @@               pure $ sizeFromName (qualName v) $ srclocOf usage     onDim _ d = pure d +-- | How the loop body produces the size at a given loop parameter position,+-- relative to the fresh variable standing for that position's size. See Note+-- [Loop size inference].+data DimClass+  = -- | A genuinely new size: this position is variant.+    Fresh+  | -- | The body produces this position's initial size (the size the fresh+    -- variable replaced). Either it produces that size directly (@Nothing@), or+    -- it copies the /current/ size of another parameter @u@ that started at the+    -- same initial size (@Just u@). In the latter case the two sizes coincide+    -- only while @u@ is unchanged, so this position is invariant exactly when+    -- @u@ is.+    Reproduces (Maybe VName)+ data ArgSource = Initial | BodyResult  wellTypedLoopArg :: ArgSource -> [VName] -> Pat ParamType -> Exp -> TermTypeM ()@@ -102,7 +115,7 @@  -- | An un-checked loop. type UncheckedLoop =-  (PatBase NoInfo VName ParamType, LoopInitBase NoInfo VName, LoopFormBase NoInfo VName, ExpBase NoInfo VName)+  (Pat ParamType, LoopInitBase Info VName, LoopFormBase Info VName, Exp)  -- | A loop that has been type-checked. type CheckedLoop =@@ -151,46 +164,16 @@     [] -> pure ()  -- | Type-check a @loop@ expression, passing in a function for--- type-checking subexpressions.+-- type-checking subexpressions.  See Note [Loop size inference]. checkLoop ::-  (ExpBase NoInfo VName -> TermTypeM Exp) ->+  (Exp -> TermTypeM Exp) ->   UncheckedLoop ->   SrcLoc ->   TermTypeM (CheckedLoop, AppRes) checkLoop checkExp (looppat, loopinit, form, loopbody) loc = do-  loopinit' <- checkExp $ case loopinit of-    LoopInitExplicit e -> e-    LoopInitImplicit _ ->-      -- Should have been filled out in Names-      error "Unspected LoopInitImplicit"+  loopinit' <- checkExp $ loopInitExp loopinit   known_before <- M.keysSet <$> getConstraints-  zeroOrderType-    (mkUsage loopinit' "use as loop variable")-    "type used as loop variable"-    . toStruct-    =<< expTypeFully loopinit' -  -- The handling of dimension sizes is a bit intricate, but very-  -- similar to checking a function, followed by checking a call to-  -- it.  The overall procedure is as follows:-  ---  -- (1) All empty dimensions in the loop pattern are instantiated-  -- with nonrigid size variables.  All explicitly specified-  -- dimensions are preserved.-  ---  -- (2) The body of the loop is type-checked.  The result type is-  -- combined with the loop pattern type to determine which sizes are-  -- variant, and these are turned into size parameters for the loop-  -- pattern.-  ---  -- (3) We now conceptually have a function parameter type and-  -- return type.  We check that it can be called with the body type-  -- as argument.-  ---  -- (4) Similarly to (3), we check that the "function" can be-  -- called with the initial loop values as argument.  The result-  -- of this is the type of the loop as a whole.-   (loop_t, new_dims_map) <-     -- dim handling (1)     allDimsFreshInType@@ -228,25 +211,65 @@         -- This works because we know that each dimension from         -- new_dims in the pattern is unique and distinct.         areSameSize <- getAreSame-        let onDims _ x y+        let initialOf v = snd <$> L.find (areSameSize v . fst) new_dims_to_initial_dim+            sameSize (Var x _ _) (Var y _ _) = areSameSize (qualLeaf x) (qualLeaf y)+            sameSize x y = x == y+            -- If the body-produced size 'd' is another loop parameter that+            -- shares the initial size 'e'', return that parameter.+            sharesInitial e' (Var d _ _)+              | Just d_init <- initialOf (qualLeaf d),+                sameSize d_init e' =+                  Just $ qualLeaf d+            sharesInitial _ _ = Nothing+            -- Classify each new_dim 'v' (with initial size 'e'') by how the+            -- loop body produces the corresponding size 'd'.+            onDims _ x y               | x == y = pure x             onDims _ e d = do-              forM_ (fvVars $ freeInExp e) $ \v -> do-                case L.find (areSameSize v . fst) new_dims_to_initial_dim of-                  Just (_, e') ->-                    if e' == d-                      then modify $ first $ M.insert v $ ExpSubst e'-                      else-                        unless (v `S.member` known_before) $-                          modify (second (v :))-                  _ ->+              forM_ (fvVars $ freeInExp e) $ \v ->+                case initialOf v of+                  Just e'+                    | sameSize e' d ->+                        modify $ M.insert v $ Reproduces Nothing+                    | v `elem` new_dims,+                      Just u <- sharesInitial e' d ->+                        modify $ M.insert v $ Reproduces (Just u)+                    | not $ v `S.member` known_before ->+                        modify $ M.insert v Fresh+                    | otherwise ->+                        pure ()+                  Nothing ->                     pure ()               pure e         loopbody_t' <- normTypeFully loopbody_t         loop_t' <- normTypeFully loop_t -        let (init_substs, sparams) =+        let classified =               execState (matchDims onDims loop_t' loopbody_t') mempty+            -- The variant sizes are the least set containing every 'Fresh'+            -- position and closed under the copies-from dependency.  See Note+            -- [Loop size inference].+            seeds = S.fromList [v | (v, Fresh) <- M.toList classified]+            grow vs =+              vs+                <> S.fromList+                  [ v+                  | (v, Reproduces (Just u)) <- M.toList classified,+                    u `S.member` vs+                  ]+            fixVariant vs =+              let vs' = grow vs+               in if vs' == vs then vs else fixVariant vs'+            variant = fixVariant seeds+            -- The invariant positions are resolved back to the initial size+            -- they stand for.+            init_substs =+              M.fromList $ do+                (v, Reproduces _) <- M.toList classified+                guard $ v `S.notMember` variant+                Just e' <- [initialOf v]+                pure (v, ExpSubst e')+            sparams = S.toList variant          checkForEscaped loc initial_levels sparams @@ -279,12 +302,15 @@   (sparams, looppat', form', loopbody') <-     case form of       For i uboundexp -> do-        uboundexp' <--          require "being the bound in a 'for' loop" anySignedType-            =<< checkExp uboundexp-        bound_t <- expTypeFully uboundexp'-        bindingIdent i bound_t $ \i' ->-          bindingPat [] looppat loop_t $ \looppat' -> incLevel $ do+        uboundexp' <- checkExp uboundexp+        it <- expType uboundexp'+        let i' = i {identType = Info it}+        bindingIdent i' $+          -- Note: bindingParam, not bindingPat, because we must+          -- preserve the Diet of the pattern as determined by the+          -- unsized type checker (the loop parameter may be+          -- consumable).+          bindingParam looppat loop_t $ \looppat' -> incLevel $ do             loopbody' <- checkExp loopbody             (sparams, looppat'') <- checkLoopReturnSize looppat' loopbody'             pure@@ -294,14 +320,13 @@                 loopbody'               )       ForIn xpat e -> do-        (arr_t, _) <- newArrayType (mkUsage' (srclocOf e)) "e" 1-        e' <- unifies "being iterated in a 'for-in' loop" arr_t =<< checkExp e+        e' <- checkExp e         t <- expTypeFully e'         case t of           _             | Just t' <- peelArray 1 t ->                 bindingPat [] xpat t' $ \xpat' ->-                  bindingPat [] looppat loop_t $ \looppat' -> incLevel $ do+                  bindingParam looppat loop_t $ \looppat' -> incLevel $ do                     loopbody' <- checkExp loopbody                     (sparams, looppat'') <- checkLoopReturnSize looppat' loopbody'                     pure@@ -315,7 +340,7 @@                   "Iteratee of a for-in loop must be an array, but expression has type"                     <+> pretty t       While cond ->-        bindingPat [] looppat loop_t $ \looppat' ->+        bindingParam looppat loop_t $ \looppat' ->           incLevel $ do             cond' <-               checkExp cond@@ -343,3 +368,58 @@     ( (sparams, looppat', LoopInitExplicit loopinit', form', loopbody'),       AppRes (toStruct loopt) retext     )++-- Note [Loop size inference]+--+-- A loop+--+--   loop p = e_init (while/for ...) do e_body+--+-- behaves like defining a function @f p = e_body@ and then calling it+-- repeatedly as @f (f ... (f e_init))@. Size-checking mirrors that: for each+-- dimension of the loop parameter @p@ we must infer whether it stays fixed+-- across iterations (*invariant*, so it can name a size that is meaningful+-- outside the loop) or may change (*variant*, so it must be existentially+-- quantified in the loop's result type). This is done in four steps, tagged+-- "dim handling (1)".."(4)" in 'checkLoop':+--+-- (1) Instantiate every size in the type of @e_init@ with a fresh nonrigid+--     variable (via 'allDimsFreshInType'), giving @loop_t@ and a map from each+--     fresh variable to the initial size it replaced. These fresh variables+--     ('new_dims') are the sizes whose variance we must determine. Distinct+--     occurrences get distinct variables even when they replace the same+--     initial size, so that e.g. two parameters that both start out @[n]@ can+--     still evolve independently.+--+-- (2) Check @e_body@ with @p@ bound at type @loop_t@, then compare the body's+--     result type against the parameter type dimension-by-dimension (with+--     'matchDims' and 'onDims'). For a new_dim @v@ that replaced initial size+--     @e'@, look at the size @d@ the body produced in that same position and+--     classify @v@ ('DimClass'):+--+--       * If @d@ is @e'@ itself, then @v@ is invariant (@Reproduces Nothing@).+--+--       * If @d@ is another new_dim @u@ that shares @v@'s initial size, then+--         @v@ is invariant iff @u@ is (@Reproduces (Just u)@). This is the+--         swap/rotation case: parameters exchanging equally-sized arrays stay+--         fixed, but if an array one of them receives has been resized, they+--         become variant.+--+--       * If @d@ is a genuinely new size, then @v@ is variant (@Fresh@).+--+--     Variance can be mutual (in @loop (a,b) = ... in (b,a)@ each of @a@,@b@+--     copies the other's size), so we take the variant set to be the least set+--     that contains every @Fresh@ position and is closed under the copies-from+--     dependency. Parameters that only copy from one another, with no @Fresh@+--     seed feeding the cycle, are therefore never made variant -- a plain swap+--     of two @[n]@ arrays keeps its precise size, while resizing one of them+--     makes both existential. The variant new_dims become the loop's size+--     parameters ('sparams'); the invariant ones are substituted back to the+--     initial size they stand for.+--+-- (3) We now conceptually have a function parameter and return type.  Check,+--     as if calling that function, that the parameter type admits the body+--     result.+--+-- (4) Likewise check that it admits the initial values; this yields the loop's+--     overall (rigid) result type.
src/Language/Futhark/TypeChecker/Terms/Monad.hs view
@@ -23,17 +23,18 @@     expType,     expTypeFully,     constrain,-    newArrayType,-    allDimsFreshInType,+    instTyVars,+    replaceTyVars,+    replaceTyVarsAbsorbable,     updateTypes,     Names,      -- * Primitive checking     unifies,-    require,     checkTypeExpNonrigid,     lookupVar,     lookupMod,+    lookupAbsTy,      -- * Sizes     isInt64,@@ -50,8 +51,9 @@ import Control.Monad.Except import Control.Monad.Reader import Control.Monad.State.Strict+import Data.Bifunctor import Data.Bitraversable-import Data.Char (isAscii)+import Data.Foldable import Data.Map.Strict qualified as M import Data.Maybe import Data.Set qualified as S@@ -61,44 +63,33 @@ import Futhark.Util.Pretty hiding (space) import Language.Futhark import Language.Futhark.Traversals-import Language.Futhark.TypeChecker.Monad hiding (BoundV, lookupMod, stateNameSource)+import Language.Futhark.TypeChecker.Constraints (TyVar)+import Language.Futhark.TypeChecker.Error+import Language.Futhark.TypeChecker.Monad hiding (BoundV, lookupAbsTy, lookupMod, stateNameSource) import Language.Futhark.TypeChecker.Monad qualified as TypeM+import Language.Futhark.TypeChecker.Terms.Scope hiding (envToTermScope, initialTermScope, lookupQualNameEnv)+import Language.Futhark.TypeChecker.Terms.Scope qualified as Scope import Language.Futhark.TypeChecker.Types import Language.Futhark.TypeChecker.Unify-import Prelude hiding (mod)+import Prelude hiding (abs, mod)  type Names = S.Set VName -data ValBinding-  = BoundV [TypeParam] StructType-  | OverloadedF [PrimType] [Maybe PrimType] (Maybe PrimType)-  | EqualityF-  deriving (Show)- unusedSize :: (MonadTypeChecker m) => SizeBinder VName -> m a unusedSize p =   typeError p mempty . withIndexLink "unused-size" $     "Size" <+> pretty p <+> "unused in pattern." -data Inferred t-  = NoneInferred-  | Ascribed t--instance Functor Inferred where-  fmap _ NoneInferred = NoneInferred-  fmap f (Ascribed t) = Ascribed (f t)- data Checking   = CheckingApply (Maybe (QualName VName)) Exp StructType StructType   | CheckingReturn ResType StructType   | CheckingAscription StructType StructType   | CheckingLetGeneralise Name   | CheckingParams (Maybe Name)-  | CheckingPat (PatBase NoInfo VName StructType) (Inferred StructType)+  | CheckingPat (PatBase Info VName StructType) (Inferred StructType)   | CheckingLoopBody StructType StructType   | CheckingLoopInitial StructType StructType   | CheckingRecordUpdate [Name] StructType StructType-  | CheckingRequired [StructType] StructType   | CheckingBranches StructType StructType  instance Pretty Checking where@@ -168,20 +159,6 @@         <+> align (pretty actual)     where       fs' = mconcat $ punctuate "." $ map pretty fs-  pretty (CheckingRequired [expected] actual) =-    "Expression must have type"-      <+> pretty expected-      <> "."-        </> "Actual type:"-        <+> align (pretty actual)-  pretty (CheckingRequired expected actual) =-    "Type of expression must be one of "-      <+> expected'-      <> "."-        </> "Actual type:"-        <+> align (pretty actual)-    where-      expected' = commasep (map pretty expected)   pretty (CheckingBranches t1 t2) =     "Branches differ in type."       </> "Former:"@@ -193,35 +170,20 @@ -- 'TermScope' will be extended during type-checking as bindings come into -- scope. data TermEnv = TermEnv-  { termScope :: TermScope,+  { termScope :: TermScope Size,     termChecking :: Maybe Checking,     termLevel :: Level,-    termChecker :: ExpBase NoInfo VName -> TermTypeM Exp,+    termCheckExp :: ExpBase Info VName -> TermTypeM Exp,     termOuterEnv :: Env,+    termTySet :: TySet,+    termTyVars :: M.Map TyVar (TypeBase () NoUniqueness),     termImportName :: ImportName   } -data TermScope = TermScope-  { scopeVtable :: M.Map VName ValBinding,-    scopeTypeTable :: M.Map VName TypeBinding,-    scopeModTable :: M.Map VName Mod-  }-  deriving (Show)--instance Semigroup TermScope where-  TermScope vt1 tt1 mt1 <> TermScope vt2 tt2 mt2 =-    TermScope (vt2 `M.union` vt1) (tt2 `M.union` tt1) (mt1 `M.union` mt2)--envToTermScope :: Env -> TermScope-envToTermScope env =-  TermScope-    { scopeVtable = vtable,-      scopeTypeTable = envTypeTable env,-      scopeModTable = envModTable env-    }-  where-    vtable = M.map valBinding $ envVtable env-    valBinding (TypeM.BoundV tps v) = BoundV tps v+-- | The scope, with sized types. See+-- "Language.Futhark.TypeChecker.Terms.Scope".+envToTermScope :: Env -> TermScope Size+envToTermScope = Scope.envToTermScope id  withEnv :: TermEnv -> Env -> TermEnv withEnv tenv env = tenv {termScope = termScope tenv <> envToTermScope env}@@ -298,12 +260,6 @@   getConstraints = gets stateConstraints   putConstraints x = modify $ \s -> s {stateConstraints = x} -  newTypeVar loc desc = do-    i <- incCounter-    v <- newID $ mkTypeVarName desc i-    constrain v $ NoConstraint Lifted $ mkUsage' loc-    pure $ Scalar $ TypeVar mempty (qualName v) []-   curLevel = asks termLevel    newDimVar usage rigidity name = do@@ -345,57 +301,224 @@           </> indent 2 (pretty t2)           </> "do not match." --- | Instantiate a type scheme with fresh type variables for its type--- parameters. Returns the names of the fresh type variables, the--- instance list, and the instantiated type.-instantiateTypeScheme ::+-- | Register the named parameters of arrows within a type as size parameters+-- ('ParamSize'), so that unification can reconstruct dependent function types+-- by linking instantiated sizes to them (the connection between a binder and+-- its uses is erased by the unsized pass). See Note [Size Inference] in+-- Language.Futhark.TypeChecker.Terms.+registerBinders :: Loc -> TypeBase Size u -> TermTypeM ()+registerBinders loc (Scalar (Arrow _ pn _ ta (RetType _ tr))) = do+  case pn of+    Named pv -> constrain pv $ ParamSize loc+    Unnamed -> pure ()+  registerBinders loc ta+  registerBinders loc tr+registerBinders loc (Scalar (Record fs)) =+  mapM_ (registerBinders loc) fs+registerBinders loc (Scalar (Sum cs)) =+  mapM_ (mapM_ (registerBinders loc)) cs+registerBinders loc (Scalar (TypeVar _ _ targs)) =+  mapM_ onTArg targs+  where+    onTArg (TypeArgType ta) = registerBinders loc ta+    onTArg TypeArgDim {} = pure ()+registerBinders _ (Scalar Prim {}) = pure ()+registerBinders loc (Array _ _ et) = registerBinders loc (Scalar et)++-- | Replace type variables inferred by the unsized type checker+-- with their solutions, instantiating their sizes with fresh+-- (non-absorbable) size variables. See Note [Size Inference] in+-- Language.Futhark.TypeChecker.Terms.+replaceTyVars :: SrcLoc -> TypeBase Size u -> TermTypeM (TypeBase Size u)+replaceTyVars = replaceTyVarsWith False++-- | Like 'replaceTyVars', but the fresh sizes may be determined to+-- be existential by unification, like instantiated sizes. This is+-- used for holes, which adopt whatever type the context provides.+-- See Note [Size Inference] in Language.Futhark.TypeChecker.Terms.+replaceTyVarsAbsorbable :: SrcLoc -> TypeBase Size u -> TermTypeM (TypeBase Size u)+replaceTyVarsAbsorbable = replaceTyVarsWith True++replaceTyVarsWith :: Bool -> SrcLoc -> TypeBase Size u -> TermTypeM (TypeBase Size u)+replaceTyVarsWith absorbable loc orig_t = do+  tyvars <- asks termTyVars+  let f ::+        TypeBase Size u ->+        StateT (M.Map VName (TypeBase Size NoUniqueness)) TermTypeM (TypeBase Size u)+      f (Scalar (Prim t)) = pure $ Scalar $ Prim t+      f+        (Scalar (TypeVar u (QualName [] v) []))+          | Just t <- M.lookup v tyvars = do+              -- Multiple occurrences of the same type variable must+              -- be given the same sizes.+              seen <- get+              case M.lookup v seen of+                Just t' -> pure $ second (const u) t'+                Nothing -> do+                  let usage = mkUsage loc "replaceTyVars"+                  (t', drepl) <-+                    lift $ allDimsFreshInType usage Nonrigid "dv" (second (const u) t)+                  -- The sizes are instantiated sizes: 'Unlifted' unless+                  -- absorbable, so they can be linked to binders of the+                  -- type itself (reconstructing dependent function+                  -- types) but only absorb existentials when absorbable.+                  -- See Note [Size Inference].+                  lift . forM_ (M.keys drepl) $ \d ->+                    constrain d $ InstSize (if absorbable then Lifted else Unlifted) usage+                  lift $ registerBinders (locOf loc) t'+                  modify $ M.insert v $ second (const NoUniqueness) t'+                  pure t'+          | otherwise =+              pure $ Scalar (TypeVar u (QualName [] v) [])+      f (Scalar (TypeVar u qn targs)) =+        Scalar . TypeVar u qn <$> mapM onTyArg targs+        where+          onTyArg (TypeArgDim e) = pure $ TypeArgDim e+          onTyArg (TypeArgType t) = TypeArgType <$> f t+      f (Scalar (Record fs)) =+        Scalar . Record <$> traverse f fs+      f (Scalar (Sum fs)) =+        Scalar . Sum <$> traverse (mapM f) fs+      f (Scalar (Arrow u pname d ta (RetType ext tr))) = do+        ta' <- f ta+        tr' <- f tr+        pure $ Scalar $ Arrow u pname d ta' $ RetType ext tr'+      f (Array u shape t) =+        arrayOfWithAliases u shape <$> f (Scalar t)++  evalStateT (f orig_t) mempty++-- | Instantiate the type parameters of a type scheme with the types+-- inferred by the unsized type checker, creating fresh variables for+-- their sizes. See Note [Size Inference] in+-- Language.Futhark.TypeChecker.Terms.+instTyVars ::+  (Substitutable (TypeBase Size u)) =>+  SrcLoc ->+  -- | The type parameters being instantiated, along with their+  -- liftedness.+  M.Map VName Liftedness ->+  TypeBase () u ->+  TypeBase Size u ->+  TermTypeM (TypeBase Size u)+instTyVars loc names orig_t1 orig_t2 = do+  tyvars <- asks termTyVars+  let f ::+        TypeBase d u ->+        TypeBase Size u ->+        StateT (M.Map VName (TypeBase Size NoUniqueness)) TermTypeM (TypeBase Size u)+      f+        (Scalar (TypeVar u (QualName [] v1) []))+        t2+          | Just t <- M.lookup v1 tyvars =+              f (second (const u) t) t2+      f (Scalar (Record fs1)) (Scalar (Record fs2)) =+        Scalar . Record <$> sequence (M.intersectionWith f fs1 fs2)+      f (Scalar (Sum fs1)) (Scalar (Sum fs2)) =+        Scalar . Sum <$> sequence (M.intersectionWith (zipWithM f) fs1 fs2)+      -- Note: uniqueness annotations are always taken from the+      -- second type, as the first (inferred) type comes from the+      -- unsized type checker, which does not track uniqueness.+      f+        (Scalar (Arrow _ _ _ t1a (RetType _ t1r)))+        (Scalar (Arrow u pname d t2a (RetType ext t2r))) = do+          ta <- f t1a t2a+          tr <- f t1r t2r+          pure $ Scalar $ Arrow u pname d ta $ RetType ext tr+      f+        (Array _ (Shape (_ : ds1)) t1)+        (Array u (Shape (d : ds2)) t2) =+          arrayOfWithAliases u (Shape [d])+            <$> f (arrayOf (Shape ds1) (Scalar t1)) (arrayOf (Shape ds2) (Scalar t2))+      f+        (Scalar (TypeVar _ v1 targs1))+        (Scalar (TypeVar u v2 targs2))+          -- If v2 is a type parameter being instantiated, it must be+          -- handled by the general case below.+          | qualLeaf v2 `M.notMember` names,+            length targs1 == length targs2 =+              Scalar . TypeVar u v1 <$> zipWithM g targs1 targs2+          where+            g (TypeArgType t1) (TypeArgType t2) =+              TypeArgType <$> f t1 t2+            g _ targ = pure targ+      f t1 t2 = do+        let usage = mkUsage loc "instantiation"+            mkNew = fst <$> lift (allDimsFreshInType usage Nonrigid "dv" t1)+        case t2 of+          Scalar (TypeVar u (QualName [] v2) [])+            | Just l <- M.lookup v2 names -> do+                seen <- get+                case M.lookup v2 seen of+                  Nothing -> do+                    (t, drepl) <- lift $ allDimsFreshInType usage Nonrigid "dv" t1+                    -- These are canonical instantiated sizes, which+                    -- unification may determine to be existential.+                    lift $ forM_ (M.keys drepl) $ \d ->+                      constrain d $ InstSize l usage+                    -- Named parameters of arrows inside the instantiated type+                    -- are registered as size parameters, such that unification+                    -- can reconstruct dependent function types by linking+                    -- instantiated sizes to them.+                    unless (null drepl) $ lift $ registerBinders (locOf loc) t+                    modify $ M.insert v2 $ second (const NoUniqueness) t+                    pure t+                  Just t -> do+                    -- Another occurrence of an already instantiated+                    -- type parameter. The sizes must be given+                    -- distinct names, as each occurrence denotes a+                    -- distinct existential size if the instantiated+                    -- size turns out to be existential.+                    occ <- lift incCounter+                    let onDim (Var (QualName _ c) info dloc) = do+                          d <- lift $ newDimVar usage Nonrigid "dv"+                          lift $ constrain d $ CopySize c occ usage+                          pure $ Var (qualName d) info dloc+                        onDim d = pure d+                    second (const u) <$> bitraverse onDim pure t+          _ -> mkNew++  (t, seen) <- runStateT (f orig_t1 orig_t2) mempty+  -- The walk above replaces occurrences of the type parameters at the type+  -- level, but the types of expressions that occur as *sizes* in the original+  -- type may also refer to the type parameters (e.g. a size 'length s - k'+  -- where 's' has a parametric type). An ordinary substitution fixes those, as+  -- substitution on expressions descends into their types.+  let substs = M.map (Subst [] . RetType []) seen+  pure $ applySubst (`M.lookup` substs) t++-- | Instantiate a type scheme with fresh variables for its size and+-- type parameters. Returns the names of the fresh size and type+-- variables and the instantiated type.+instTypeScheme ::   QualName VName ->   SrcLoc ->   [TypeParam] ->   StructType ->+  TypeBase () NoUniqueness ->   TermTypeM ([VName], StructType)-instantiateTypeScheme qn loc tparams t = do-  let tnames = map typeParamName 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')+instTypeScheme qn loc tparams scheme_t inferred = do+  (names, substs) <- fmap (unzip . catMaybes) . forM tparams $ \tparam -> do+    case tparam of+      TypeParamType {} -> pure Nothing+      TypeParamDim v _ -> do+        i <- incCounter+        v' <- newID $ mkTypeVarName (baseName v) i+        -- The instantiation of a size parameter may turn out to be+        -- an existential size, when the value whose type contains it+        -- is returned from a function argument.+        constrain v' . InstSize Lifted . mkUsage loc . docText $+          "instantiated size parameter of " <> dquotes (pretty qn)+        pure $ Just (v', (v, ExpSubst $ sizeFromName (qualName v') loc)) --- | Create a new type name and insert it (unconstrained) in the--- substitution map.-instantiateTypeParam ::-  (Monoid as) =>-  QualName VName ->-  SrcLoc ->-  TypeParam ->-  TermTypeM (VName, Subst (RetTypeBase dim as))-instantiateTypeParam qn loc tparam = do-  i <- incCounter-  let name = nameFromText (T.takeWhile isAscii (baseText (typeParamName tparam)))-  v <- newID $ mkTypeVarName name i-  case tparam of-    TypeParamType x _ _ -> do-      constrain v . NoConstraint x . mkUsage loc . docText $-        "instantiated type parameter of " <> dquotes (pretty qn)-      pure (v, Subst [] $ RetType [] $ Scalar $ TypeVar mempty (qualName v) [])-    TypeParamDim {} -> do-      constrain v . Size Nothing . mkUsage loc . docText $-        "instantiated size parameter of " <> dquotes (pretty qn)-      pure (v, ExpSubst $ sizeFromName (qualName v) loc)+  let tp_names = M.fromList $ mapMaybe tpName tparams+      tpName (TypeParamType l v _) = Just (v, l)+      tpName TypeParamDim {} = Nothing+  t' <- instTyVars loc tp_names inferred $ applySubst (`lookup` substs) scheme_t+  pure (names, t') -lookupQualNameEnv :: QualName VName -> TermTypeM TermScope-lookupQualNameEnv (QualName [q] _)-  | isIntrinsic q = asks termScope -- Magical intrinsic module.-lookupQualNameEnv qn@(QualName quals _) = do-  scope <- asks termScope-  descend scope quals-  where-    descend scope [] = pure scope-    descend scope (q : qs)-      | Just (ModEnv q_scope) <- M.lookup q $ scopeModTable scope =-          descend (envToTermScope q_scope) qs-      | otherwise =-          error $ "lookupQualNameEnv " <> show qn+lookupQualNameEnv :: QualName VName -> TermTypeM (TermScope Size)+lookupQualNameEnv qn = asks $ \tenv -> Scope.lookupQualNameEnv id (termScope tenv) qn  lookupMod :: QualName VName -> TermTypeM Mod lookupMod qn@(QualName _ name) = do@@ -404,7 +527,7 @@     Nothing -> error $ "lookupMod: " <> show qn     Just m -> pure m -localScope :: (TermScope -> TermScope) -> TermTypeM a -> TermTypeM a+localScope :: (TermScope Size -> TermScope Size) -> TermTypeM a -> TermTypeM a localScope f = local $ \tenv -> tenv {termScope = f $ termScope tenv}  instance MonadTypeChecker TermTypeM where@@ -446,42 +569,46 @@       Nothing ->         throwError $ TypeError (locOf loc) notes s -lookupVar :: SrcLoc -> QualName VName -> TermTypeM StructType-lookupVar loc qn@(QualName qs name) = do+lookupVar :: SrcLoc -> QualName VName -> StructType -> TermTypeM StructType+lookupVar loc qn@(QualName qs name) inst_t = do   scope <- lookupQualNameEnv qn-  let usage = mkUsage loc $ docText $ "use of " <> dquotes (pretty qn)--  case M.lookup name $ scopeVtable scope of+  outer_env <- asks termOuterEnv+  -- Top-level value bindings are not in the term scope (see+  -- 'envToTermScopeNoVals'); look them up on demand in the outer+  -- environment.+  case M.lookup name (scopeVtable scope)+    `mplus` Scope.lookupOuterVal id outer_env name of     Nothing ->       error $ "lookupVar: " <> show qn-    Just (BoundV tparams t) -> do+    Just (BoundV tparams bound_t) ->       if null tparams && null qs-        then pure t+        then pure bound_t         else do-          (tnames, t') <- instantiateTypeScheme qn loc tparams t-          outer_env <- asks termOuterEnv-          pure $ qualifyTypeVars outer_env tnames qs t'-    Just EqualityF -> do-      argtype <- newTypeVar loc "t"-      equalityType usage argtype-      pure $-        Scalar . Arrow mempty Unnamed Observe argtype . RetType [] $-          Scalar $-            Arrow mempty Unnamed Observe argtype $-              RetType [] $-                Scalar $-                  Prim Bool-    Just (OverloadedF ts pts rt) -> do-      argtype <- newTypeVar loc "t"-      mustBeOneOf ts usage argtype-      let (pts', rt') = instOverloaded argtype pts rt-      pure $ foldFunType (map (toParam Observe) pts') $ RetType [] $ toRes Nonunique rt'-  where-    instOverloaded argtype pts rt =-      ( map (maybe (toStruct argtype) (Scalar . Prim)) pts,-        maybe (toStruct argtype) (Scalar . Prim) rt-      )+          (tnames, t) <- instTypeScheme qn loc tparams bound_t $ first (const ()) inst_t+          pure $ qualifyTypeVars outer_env tnames qs t+    Just EqualityF ->+      replaceTyVars loc inst_t+    Just OverloadedF {} ->+      replaceTyVars loc inst_t+    -- See Note [Checking recursive functions] in+    -- Language.Futhark.TypeChecker.Terms.+    Just RecursiveV ->+      replaceTyVars loc inst_t +-- | Look up the liftedness of an abstract type.+lookupAbsTy :: QualName VName -> TermTypeM Liftedness+lookupAbsTy v | isIntrinsic (qualLeaf v) = pure Unlifted+lookupAbsTy v = do+  abs <- asks termTySet+  case M.lookup v abs of+    Just l -> pure l+    Nothing ->+      error $+        unlines+          [ "lookupAbsTy: " <> prettyString v,+            "known: " <> show abs+          ]+ onFailure :: Checking -> TermTypeM a -> TermTypeM a onFailure c = local $ \env -> env {termChecking = Just c} @@ -501,11 +628,10 @@ incLevel :: TermTypeM a -> TermTypeM a incLevel = local $ \env -> env {termLevel = termLevel env + 1} --- | Get the type of an expression, with top level type variables--- substituted.  Never call 'typeOf' directly (except in a few--- carefully inspected locations)!+-- | Get the type of an expression. Currently no different from+-- 'typeOf', but kept as the counterpart of 'expTypeFully'. expType :: Exp -> TermTypeM StructType-expType = normType . typeOf+expType = pure . typeOf  -- | Get the type of an expression, with all type variables -- substituted.  Slower than 'expType', but sometimes necessary.@@ -514,33 +640,6 @@ expTypeFully :: Exp -> TermTypeM StructType expTypeFully = normTypeFully . typeOf -newArrayType :: Usage -> Name -> Int -> TermTypeM (StructType, StructType)-newArrayType usage desc r = do-  v <- newTypeName desc-  constrain v $ NoConstraint Unlifted usage-  dims <- replicateM r $ newDimVar usage Nonrigid "dim"-  let rowt = TypeVar mempty (qualName v) []-      mkSize = flip sizeFromName (srclocOf usage) . qualName-  pure-    ( Array mempty (Shape $ map mkSize dims) rowt,-      Scalar rowt-    )---- | Replace *all* dimensions with distinct fresh size variables.-allDimsFreshInType ::-  Usage ->-  Rigidity ->-  Name ->-  TypeBase Size als ->-  TermTypeM (TypeBase Size als, M.Map VName Size)-allDimsFreshInType usage r desc t =-  runStateT (bitraverse onDim pure t) mempty-  where-    onDim d = do-      v <- lift $ newDimVar usage r desc-      modify $ M.insert v d-      pure $ sizeFromName (qualName v) $ srclocOf usage- -- | Replace all type variables with their concrete types. updateTypes :: (ASTMappable e) => e -> TermTypeM e updateTypes = astMap tv@@ -561,24 +660,15 @@   unify (mkUsage (srclocOf e) why) t . toStruct =<< expType e   pure e --- | @require ts e@ causes a 'TypeError' if @expType e@ is not one of--- the types in @ts@.  Otherwise, simply returns @e@.-require :: T.Text -> [PrimType] -> Exp -> TermTypeM Exp-require why ts e = do-  mustBeOneOf ts (mkUsage (srclocOf e) why) . toStruct =<< expType e-  pure e--checkExpForSize :: ExpBase NoInfo VName -> TermTypeM Exp+checkExpForSize :: ExpBase Info VName -> TermTypeM Exp checkExpForSize e = do-  checker <- asks termChecker+  checker <- asks termCheckExp   e' <- checker e   let t = toStruct $ typeOf e'   unify (mkUsage (locOf e') "Size expression") t (Scalar (Prim (Signed Int64)))   updateTypes e' -checkTypeExpNonrigid ::-  TypeExp (ExpBase NoInfo VName) VName ->-  TermTypeM (TypeExp Exp VName, ResType, [VName])+checkTypeExpNonrigid :: TypeExp Exp VName -> TermTypeM (TypeExp Exp VName, ResType, [VName]) checkTypeExpNonrigid te = do   (te', svars, rettype, _l) <- checkTypeExp checkExpForSize te @@ -602,50 +692,26 @@  -- Running -initialTermScope :: TermScope-initialTermScope =-  TermScope-    { scopeVtable = initialVtable,-      scopeTypeTable = mempty,-      scopeModTable = mempty-    }-  where-    initialVtable = M.fromList $ mapMaybe addIntrinsicF $ M.toList intrinsics--    prim = Scalar . Prim-    arrow x y = Scalar $ Arrow mempty Unnamed Observe x y--    addIntrinsicF (name, IntrinsicMonoFun pts t) =-      Just (name, BoundV [] $ arrow pts' $ RetType [] $ prim t)-      where-        pts' = case pts of-          [pt] -> prim pt-          _ -> Scalar $ tupleRecord $ map prim pts-    addIntrinsicF (name, IntrinsicOverloadedFun ts pts rts) =-      Just (name, OverloadedF ts pts rts)-    addIntrinsicF (name, IntrinsicPolyFun tvs pts rt) =-      Just-        ( name,-          BoundV tvs $ foldFunType pts rt-        )-    addIntrinsicF (name, IntrinsicEquality) =-      Just (name, EqualityF)-    addIntrinsicF _ = Nothing+initialTermScope :: TermScope Size+initialTermScope = Scope.initialTermScope id -runTermTypeM :: (ExpBase NoInfo VName -> TermTypeM Exp) -> TermTypeM a -> TypeM a-runTermTypeM checker (TermTypeM m) = do-  initial_scope <- (initialTermScope <>) . envToTermScope <$> askEnv+runTermTypeM :: (ExpBase Info VName -> TermTypeM Exp) -> M.Map TyVar (TypeBase () NoUniqueness) -> TermTypeM a -> TypeM a+runTermTypeM checker tyvars (TermTypeM m) = do+  initial_scope <- (initialTermScope <>) . Scope.envToTermScopeNoVals <$> askEnv   name <- askImportName   outer_env <- askEnv   src <- gets TypeM.stateNameSource+  abs <- getTySet   let initial_tenv =         TermEnv           { termScope = initial_scope,             termChecking = Nothing,             termLevel = 0,-            termChecker = checker,+            termCheckExp = checker,             termImportName = name,-            termOuterEnv = outer_env+            termOuterEnv = outer_env,+            termTySet = abs,+            termTyVars = tyvars           }       initial_state =         TermTypeState
src/Language/Futhark/TypeChecker/Terms/Pat.hs view
@@ -2,6 +2,7 @@ module Language.Futhark.TypeChecker.Terms.Pat   ( binding,     bindingParams,+    bindingParam,     bindingPat,     bindingIdent,     bindingSizes,@@ -11,7 +12,7 @@ import Control.Monad import Data.Bifunctor import Data.Either-import Data.List (find, isPrefixOf, sort, sortBy)+import Data.List (find, isPrefixOf, sortBy) import Data.Map.Strict qualified as M import Data.Maybe import Data.Ord (comparing)@@ -20,6 +21,7 @@ import Language.Futhark import Language.Futhark.TypeChecker.Monad hiding (BoundV) import Language.Futhark.TypeChecker.Terms.Monad+import Language.Futhark.TypeChecker.Terms.Scope (typeParamIdent) import Language.Futhark.TypeChecker.Types import Language.Futhark.TypeChecker.Unify hiding (Usage) import Prelude hiding (mod)@@ -78,18 +80,12 @@   binding (mapMaybe typeParamIdent tparams)     . bindingTypes (concatMap typeParamType tparams)   where-    typeParamType (TypeParamType l v loc) =-      [ Left (v, TypeAbbr l [] $ RetType [] $ Scalar (TypeVar mempty (qualName v) [])),-        Right (v, ParamType l $ locOf loc)+    typeParamType (TypeParamType l v _) =+      [ Left (v, TypeAbbr l [] $ RetType [] $ Scalar (TypeVar mempty (qualName v) []))       ]     typeParamType (TypeParamDim v loc) =       [Right (v, ParamSize $ locOf loc)] -typeParamIdent :: TypeParam -> Maybe (Ident StructType)-typeParamIdent (TypeParamDim v loc) =-  Just $ Ident v (Info $ Scalar $ Prim $ Signed Int64) loc-typeParamIdent _ = Nothing- -- | Bind @let@-bound sizes.  This is usually followed by 'bindingPat' -- immediately afterwards. bindingSizes :: [SizeBinder VName] -> TermTypeM a -> TermTypeM a@@ -100,58 +96,40 @@       Ident (sizeName size) (Info (Scalar (Prim (Signed Int64)))) (srclocOf size)  -- | Bind a single term-level identifier.-bindingIdent ::-  IdentBase NoInfo VName StructType ->-  StructType ->-  (Ident StructType -> TermTypeM a) ->-  TermTypeM a-bindingIdent (Ident v NoInfo vloc) t m = do-  let ident = Ident v (Info t) vloc-  binding [ident] $ m ident---- All this complexity is just so we can handle un-suffixed numeric--- literals in patterns.-patLitMkType :: PatLit -> SrcLoc -> TermTypeM ParamType-patLitMkType (PatLitInt _) loc = do-  t <- newTypeVar loc "t"-  mustBeOneOf anyNumberType (mkUsage loc "integer literal") (toStruct t)-  pure t-patLitMkType (PatLitFloat _) loc = do-  t <- newTypeVar loc "t"-  mustBeOneOf anyFloatType (mkUsage loc "float literal") (toStruct t)-  pure t-patLitMkType (PatLitPrim v) _ =-  pure $ Scalar $ Prim $ primValueType v+bindingIdent :: Ident StructType -> TermTypeM a -> TermTypeM a+bindingIdent ident = binding [ident]  checkPat' ::   [(SizeBinder VName, QualName VName)] ->-  PatBase NoInfo VName ParamType ->+  Pat ParamType ->   Inferred ParamType ->   TermTypeM (Pat ParamType) checkPat' sizes (PatParens p loc) t =   PatParens <$> checkPat' sizes p t <*> pure loc checkPat' sizes (PatAttr attr p loc) t =   PatAttr <$> checkAttr attr <*> checkPat' sizes p t <*> pure loc-checkPat' _ (Id name NoInfo loc) (Ascribed t) =-  pure $ Id name (Info t) loc-checkPat' _ (Id name NoInfo loc) NoneInferred = do-  t <- newTypeVar loc "t"-  pure $ Id name (Info t) loc-checkPat' _ (Wildcard _ loc) (Ascribed t) =-  pure $ Wildcard (Info t) loc-checkPat' _ (Wildcard NoInfo loc) NoneInferred = do-  t <- newTypeVar loc "t"-  pure $ Wildcard (Info t) loc-checkPat' sizes p@(TuplePat ps loc) (Ascribed t)+checkPat' _ (Id name (Info t) loc) NoneInferred = do+  t' <- replaceTyVars loc t+  pure $ Id name (Info t') loc+checkPat' _ (Id name (Info t1) loc) (Ascribed t2) = do+  t' <- instTyVars loc mempty (first (const ()) t1) t2+  pure $ Id name (Info t') loc+checkPat' _ (Wildcard (Info t) loc) NoneInferred = do+  t' <- replaceTyVars loc t+  pure $ Wildcard (Info t') loc+checkPat' _ (Wildcard (Info t1) loc) (Ascribed t2) = do+  t' <- instTyVars loc mempty (first (const ()) t1) t2+  pure $ Wildcard (Info t') loc+checkPat' sizes (TuplePat ps loc) (Ascribed t)   | Just ts <- isTupleRecord t,     length ts == length ps =       TuplePat         <$> zipWithM (checkPat' sizes) ps (map Ascribed ts)         <*> pure loc-  | otherwise = do-      ps_t <- replicateM (length ps) (newTypeVar loc "t")-      unify (mkUsage loc "matching a tuple pattern") (Scalar (tupleRecord ps_t)) (toStruct t)-      checkPat' sizes p $ Ascribed $ toParam Observe $ Scalar $ tupleRecord ps_t+  | otherwise =+      -- The unsized type checker has already verified that the+      -- pattern matches the type of the bound expression.+      error $ "checkPat' TuplePat: " <> prettyString (toStruct t) checkPat' sizes (TuplePat ps loc) NoneInferred =   TuplePat <$> mapM (\p -> checkPat' sizes p NoneInferred) ps <*> pure loc checkPat' _ (RecordPat p_fs _) _@@ -161,21 +139,16 @@           </> "Did you mean"           <> dquotes (pretty (drop 1 (nameToString f)) <> "=_")           <> "?"-checkPat' sizes p@(RecordPat p_fs loc) (Ascribed t)+checkPat' sizes (RecordPat p_fs loc) (Ascribed t)   | Scalar (Record t_fs) <- t,     p_fs' <- sortBy (comparing fst) p_fs,     t_fs' <- sortBy (comparing fst) (M.toList t_fs),     map fst t_fs' == map (unLoc . fst) p_fs' =       RecordPat <$> zipWithM check p_fs' t_fs' <*> pure loc-  | otherwise = do-      p_fs' <- traverse (const $ newTypeVar loc "t") $ M.fromList $ map (first unLoc) p_fs--      when (sort (M.keys p_fs') /= sort (map (unLoc . fst) p_fs)) $-        typeError loc mempty $-          "Duplicate fields in record pattern" <+> pretty p <> "."--      unify (mkUsage loc "matching a record pattern") (Scalar (Record p_fs')) (toStruct t)-      checkPat' sizes p $ Ascribed $ toParam Observe $ Scalar (Record p_fs')+  | otherwise =+      -- The unsized type checker has already verified that the+      -- pattern matches the type of the bound expression.+      error $ "checkPat' RecordPat: " <> prettyString (toStruct t)   where     check (L f_loc f, p_f) (_, t_f) = (L f_loc f,) <$> checkPat' sizes p_f (Ascribed t_f) checkPat' sizes (RecordPat fs loc) NoneInferred =@@ -199,54 +172,33 @@         <$> checkPat' sizes p (Ascribed (resToParam st))         <*> pure t'         <*> pure loc-checkPat' _ (PatLit l NoInfo loc) (Ascribed t) = do-  t' <- patLitMkType l loc-  unify (mkUsage loc "matching against literal") (toStruct t') (toStruct t)-  pure $ PatLit l (Info t') loc-checkPat' _ (PatLit l NoInfo loc) NoneInferred = do-  t' <- patLitMkType l loc+checkPat' _ (PatLit l (Info t) loc) _ = do+  t' <- replaceTyVars loc t   pure $ PatLit l (Info t') loc-checkPat' sizes (PatConstr n NoInfo ps loc) (Ascribed (Scalar (Sum cs)))+checkPat' sizes (PatConstr n info ps loc) NoneInferred = do+  ps' <- mapM (\p -> checkPat' sizes p NoneInferred) ps+  pure $ PatConstr n info ps' loc+checkPat' sizes (PatConstr n _ ps loc) (Ascribed (Scalar (Sum cs)))   | Just ts <- M.lookup n cs = do-      when (length ps /= length ts) $-        typeError loc mempty $-          "Pattern #"-            <> pretty n-            <> " expects"-              <+> pretty (length ps)-              <+> "constructor arguments, but type provides"-              <+> pretty (length ts)-              <+> "arguments."-      ps' <- zipWithM (checkPat' sizes) ps $ map Ascribed ts+      ps' <- zipWithM (\p t -> checkPat' sizes p (Ascribed t)) ps ts       pure $ PatConstr n (Info (Scalar (Sum cs))) ps' loc-checkPat' sizes (PatConstr n NoInfo ps loc) (Ascribed t) = do-  t' <- newTypeVar loc "t"-  ps' <- forM ps $ \p -> do-    p_t <- newTypeVar (srclocOf p) "t"-    checkPat' sizes p $ Ascribed p_t-  mustHaveConstr usage n (toStruct t') (patternStructType <$> ps')-  unify usage t' (toStruct t)-  pure $ PatConstr n (Info t) ps' loc-  where-    usage = mkUsage loc "matching against constructor"-checkPat' sizes (PatConstr n NoInfo ps loc) NoneInferred = do-  ps' <- mapM (\p -> checkPat' sizes p NoneInferred) ps-  t <- newTypeVar loc "t"-  mustHaveConstr usage n (toStruct t) (patternStructType <$> ps')-  pure $ PatConstr n (Info t) ps' loc-  where-    usage = mkUsage loc "matching against constructor"+checkPat' _ p t =+  error . unlines $+    [ "checkPat': bad case",+      prettyString p,+      show t+    ]  checkPat ::   [(SizeBinder VName, QualName VName)] ->-  PatBase NoInfo VName (TypeBase Size u) ->+  Pat ParamType ->   Inferred StructType ->   (Pat ParamType -> TermTypeM a) ->   TermTypeM a checkPat sizes p t m = do   p' <-     onFailure (CheckingPat (fmap toStruct p) t) $-      checkPat' sizes (fmap (toParam Observe) p) (fmap (toParam Observe) t)+      checkPat' sizes p (fmap (toParam Observe) t)    let explicit = mustBeExplicitInType $ patternStructType p' @@ -259,19 +211,30 @@     [] ->       m p' +-- | Check and bind a single parameter.+bindingParam ::+  Pat ParamType ->+  StructType ->+  (Pat ParamType -> TermTypeM a) ->+  TermTypeM a+bindingParam p t m = do+  checkPat mempty p (Ascribed t) $ \p' ->+    binding (patIdents (fmap toStruct p')) $ m p'+ -- | Check and bind a @let@-pattern. bindingPat ::   [SizeBinder VName] ->-  PatBase NoInfo VName (TypeBase Size u) ->+  Pat (TypeBase Size u) ->   StructType ->   (Pat ParamType -> TermTypeM a) ->   TermTypeM a bindingPat sizes p t m = do   substs <- mapM mkSizeSubst sizes-  checkPat substs p (Ascribed t) $ \p' -> binding (patIdents (fmap toStruct p')) $-    case filter ((`S.notMember` fvVars (freeInPat p')) . sizeName) sizes of-      [] -> m p'-      size : _ -> unusedSize size+  checkPat substs (fmap (toParam Observe) p) (Ascribed t) $ \p' ->+    binding (patIdents (fmap toStruct p')) $+      case filter ((`S.notMember` fvVars (freeInPat p')) . sizeName) sizes of+        [] -> m p'+        size : _ -> unusedSize size   where     mkSizeSubst v = do       v' <- newID $ baseName $ sizeName v@@ -282,13 +245,15 @@ -- | Check and bind type and value parameters. bindingParams ::   [TypeParam] ->-  [PatBase NoInfo VName ParamType] ->+  [Pat ParamType] ->   ([Pat ParamType] -> TermTypeM a) ->   TermTypeM a bindingParams tps orig_ps m = bindingTypeParams tps $ do   let descend ps' (p : ps) =         checkPat [] p NoneInferred $ \p' ->-          binding (patIdents $ fmap toStruct p') $ incLevel $ descend (p' : ps') ps+          binding (patIdents $ fmap toStruct p') $+            incLevel $+              descend (p' : ps') ps       descend ps' [] = m $ reverse ps'    incLevel $ descend [] orig_ps
+ src/Language/Futhark/TypeChecker/Terms/Scope.hs view
@@ -0,0 +1,164 @@+-- | The notion of scope used during term checking, and other small+-- definitions shared between the unsized term checker+-- ("Language.Futhark.TypeChecker.Terms.Unsized") and the sized one+-- ("Language.Futhark.TypeChecker.Terms"). The two checkers represent+-- sizes differently, so the definitions here are parameterised over+-- the size representation, and the functions take a function for+-- converting the sizes of module-level types.+module Language.Futhark.TypeChecker.Terms.Scope+  ( ValBinding (..),+    TermScope (..),+    Inferred (..),+    envToTermScope,+    envToTermScopeNoVals,+    lookupOuterVal,+    initialTermScope,+    lookupQualNameEnv,+    typeParamIdent,+  )+where++import Data.Map.Strict qualified as M+import Data.Maybe+import Language.Futhark+import Language.Futhark.TypeChecker.Monad hiding (BoundV)+import Language.Futhark.TypeChecker.Monad qualified as TypeM++-- | What a bound value stands for. Note that although modules are in the same+-- name space, they are not value bindings.+data ValBinding dim+  = BoundV [TypeParam] (TypeBase dim NoUniqueness)+  | OverloadedF [PrimType] [Maybe PrimType] (Maybe PrimType)+  | EqualityF+  | -- | A recursive function with no declared return type, currently being+    -- checked and in scope of its own body. Each occurrence adopts the type+    -- recorded for it by the unsized type checker, with fresh (unrelated) sizes.+    -- See Note [Checking recursive functions] in+    -- Language.Futhark.TypeChecker.Terms.+    RecursiveV+  deriving (Show)++-- | The lexical scope information.+data TermScope dim = TermScope+  { scopeVtable :: M.Map VName (ValBinding dim),+    scopeTypeTable :: M.Map VName TypeBinding,+    scopeModTable :: M.Map VName Mod+  }+  deriving (Show)++instance Semigroup (TermScope dim) where+  TermScope vt1 tt1 mt1 <> TermScope vt2 tt2 mt2 =+    TermScope (vt2 `M.union` vt1) (tt2 `M.union` tt1) (mt1 `M.union` mt2)++-- | During type checking a pattern, we might find an explicit ascription. These+-- contain complete type information (although they must of course still be+-- checked against what remains of the pattern).+data Inferred t+  = NoneInferred+  | Ascribed t+  deriving (Show)++instance Functor Inferred where+  fmap _ NoneInferred = NoneInferred+  fmap f (Ascribed t) = Ascribed (f t)++-- | Create a scope from a module-level environment.+envToTermScope ::+  (StructType -> TypeBase dim NoUniqueness) ->+  Env ->+  TermScope dim+envToTermScope onType env =+  TermScope+    { scopeVtable = vtable,+      scopeTypeTable = envTypeTable env,+      scopeModTable = envModTable env+    }+  where+    vtable = M.map valBinding $ envVtable env+    valBinding (TypeM.BoundV tps v) = BoundV tps $ onType v++-- | Like 'envToTermScope', but omits the (potentially very large)+-- value table. Its bindings are instead looked up on demand with+-- 'lookupOuterVal'. This avoids transforming the entire value table of+-- the outer environment - which for a top-level definition includes+-- the whole prelude - once for every binding we check.+envToTermScopeNoVals :: Env -> TermScope dim+envToTermScopeNoVals env =+  TermScope+    { scopeVtable = mempty,+      scopeTypeTable = envTypeTable env,+      scopeModTable = envModTable env+    }++-- | Look up a single value binding in an 'Env' and convert it, using+-- the given size conversion. The fallback for names not found in the+-- (value-free) term scope built by 'envToTermScopeNoVals'.+lookupOuterVal ::+  (StructType -> TypeBase dim NoUniqueness) ->+  Env ->+  VName ->+  Maybe (ValBinding dim)+lookupOuterVal onType env v =+  convert <$> M.lookup v (envVtable env)+  where+    convert (TypeM.BoundV tps t) = BoundV tps $ onType t++-- | The initial scope, containing the intrinsics.+initialTermScope ::+  (StructType -> TypeBase dim NoUniqueness) ->+  TermScope dim+initialTermScope onType =+  TermScope+    { scopeVtable = initialVtable,+      scopeTypeTable = mempty,+      scopeModTable = mempty+    }+  where+    initialVtable = M.fromList $ mapMaybe addIntrinsicF $ M.toList intrinsics++    prim = Scalar . Prim+    arrow x y = Scalar $ Arrow mempty Unnamed Observe x y++    addIntrinsicF (name, IntrinsicMonoFun pts t) =+      Just (name, BoundV [] $ onType $ arrow pts' $ RetType [] $ prim t)+      where+        pts' = case pts of+          [pt] -> prim pt+          _ -> Scalar $ tupleRecord $ map prim pts+    addIntrinsicF (name, IntrinsicOverloadedFun ts pts rts) =+      Just (name, OverloadedF ts pts rts)+    addIntrinsicF (name, IntrinsicPolyFun tvs pts rt) =+      Just+        ( name,+          BoundV tvs $ onType $ foldFunType pts rt+        )+    addIntrinsicF (name, IntrinsicEquality) =+      Just (name, EqualityF)+    addIntrinsicF _ = Nothing++-- | Find the scope corresponding to the qualifiers of the given+-- name. Fails with 'error' if the qualifiers do not name a module,+-- as this means the program should not have made it through earlier+-- checks.+lookupQualNameEnv ::+  (StructType -> TypeBase dim NoUniqueness) ->+  TermScope dim ->+  QualName VName ->+  TermScope dim+lookupQualNameEnv _ scope (QualName [q] _)+  | isIntrinsic q = scope -- Magical intrinsic module.+lookupQualNameEnv onType scope qn@(QualName quals _) = descend scope quals+  where+    descend s [] = s+    descend s (q : qs)+      | Just (ModEnv q_env) <- M.lookup q $ scopeModTable s =+          descend (envToTermScope onType q_env) qs+      | otherwise =+          error $ "lookupQualNameEnv " <> show qn++-- | An identifier corresponding to a type parameter, for size+-- parameters, which also exist as terms.+typeParamIdent :: TypeParam -> Maybe (Ident StructType)+typeParamIdent (TypeParamDim v loc) =+  Just $ Ident v (Info $ Scalar $ Prim $ Signed Int64) loc+typeParamIdent _ = Nothing
+ src/Language/Futhark/TypeChecker/Terms/Unsized.hs view
@@ -0,0 +1,1406 @@+{-# LANGUAGE Strict #-}++-- | Unsized type checking.+--+-- This checker generates type constraints (type 'CtTy') which are then solved+-- to find a solution. The result is a decorated AST where most of the type+-- annotations are just references to type variables. Further, all the+-- size-specific annotations (e.g. existential sizes) just contain dummy values,+-- such as empty lists.+--+-- If Futhark had no fancy type system features, then this pass would+-- essentially be all you needed.+module Language.Futhark.TypeChecker.Terms.Unsized+  ( checkValDef,+    checkSingleExp,+    checkSizeExp,+    Solution,+  )+where++import Control.Monad+import Control.Monad.Except+import Control.Monad.Reader+import Control.Monad.State.Strict+import Data.Bifoldable (bifoldMap)+import Data.Bifunctor+import Data.Bitraversable+import Data.Char (isAscii)+import Data.Either (partitionEithers)+import Data.List qualified as L+import Data.List.NonEmpty qualified as NE+import Data.Loc (Loc (NoLoc))+import Data.Map qualified as M+import Data.Maybe+import Data.Ord (comparing)+import Data.Set qualified as S+import Data.Text qualified as T+import Futhark.FreshNames qualified as FreshNames+import Futhark.MonadFreshNames hiding (newName)+import Futhark.Util (nubOrd)+import Futhark.Util.Pretty+import Language.Futhark+import Language.Futhark.TypeChecker.Constraints+import Language.Futhark.TypeChecker.Monad hiding (BoundV, lookupMod)+import Language.Futhark.TypeChecker.Monad qualified as TypeM+import Language.Futhark.TypeChecker.Terms.Scope hiding (envToTermScope, initialTermScope, lookupQualNameEnv)+import Language.Futhark.TypeChecker.Terms.Scope qualified as Scope+import Language.Futhark.TypeChecker.TySolve hiding (Type)+import Language.Futhark.TypeChecker.Types+import Language.Futhark.TypeChecker.Unify (mkUsage)+import Prelude hiding (mod)++type Type = CtType ()++-- | The unsized checker ignores sizes entirely, so we erase them. (The+-- constraint solver never inspects sizes; see '()' for the shape+-- representation this could use if size-aware rank inference were ever+-- wired up.)+toType :: TypeBase Size u -> TypeBase () u+toType = first (const ())++-- | Type checking happens with access to this environment.  The+-- 'TermScope' will be extended during type-checking as bindings come into+-- scope.+data TermEnv = TermEnv+  { termScope :: TermScope (),+    termLevel :: Level,+    termOuterEnv :: Env,+    termImportName :: ImportName,+    -- | The liftedness of abstract types.+    termTySet :: TySet+  }++-- | An instantiation (at the given location, of a type parameter of+-- the given polymorphic function, with the given liftedness, as the+-- given type variable), recorded such that we can check, after+-- constraint solving, that the liftedness of the type parameter is+-- respected ('checkTyInstLiftedness').+data TyInst = TyInst Loc (QualName VName) Liftedness TyVar++-- | The state is a set of constraints and a counter for generating+-- type names.  This is distinct from the usual counter we use for+-- generating unique names, as these will be user-visible.+data TermState = TermState+  { termConstraints :: [CtTy ()],+    termTyVars :: TyVars (),+    termTyParams :: TyParams,+    termTyInsts :: [TyInst],+    termCounter :: !Int,+    termWarnings :: Warnings,+    termNameSource :: VNameSource,+    -- | Mapping from artificial type variables to the actual types they represent.+    termArtificial :: M.Map TyVar Type+  }++newtype TermM a+  = TermM+      ( ReaderT+          TermEnv+          (StateT TermState (Except (Warnings, TypeError)))+          a+      )+  deriving+    ( Monad,+      Functor,+      Applicative,+      MonadReader TermEnv,+      MonadState TermState+    )++-- | The scope, with sizes erased. See+-- "Language.Futhark.TypeChecker.Terms.Scope".+envToTermScope :: Env -> TermScope ()+envToTermScope = Scope.envToTermScope toType++initialTermScope :: TermScope ()+initialTermScope = Scope.initialTermScope toType++runTermM :: TermM a -> TypeM a+runTermM (TermM m) = do+  initial_scope <- (initialTermScope <>) . Scope.envToTermScopeNoVals <$> askEnv+  name <- askImportName+  outer_env <- askEnv+  src <- gets stateNameSource+  abs_types <- getTySet+  let initial_env =+        TermEnv+          { termScope = initial_scope,+            termLevel = 0,+            termImportName = name,+            termOuterEnv = outer_env,+            termTySet = abs_types+          }+      initial_state =+        TermState+          { termConstraints = mempty,+            termTyVars = mempty,+            termTyParams = mempty,+            termTyInsts = mempty,+            termWarnings = mempty,+            termNameSource = src,+            termCounter = 0,+            termArtificial = mempty+          }+  case runExcept (runStateT (runReaderT m initial_env) initial_state) of+    Left (ws, e) -> do+      warnings ws+      throwError e+    Right (a, TermState {termNameSource, termWarnings}) -> do+      warnings termWarnings+      modify $ \s -> s {stateNameSource = termNameSource}+      pure a++incLevel :: TermM a -> TermM a+incLevel = local $ \env -> env {termLevel = termLevel env + 1}++curLevel :: TermM Int+curLevel = asks termLevel++incCounter :: TermM Int+incCounter = do+  s <- get+  put s {termCounter = termCounter s + 1}+  pure $ termCounter s++tyVarType :: u -> TyVar -> TypeBase dim u+tyVarType u v = Scalar $ TypeVar u (qualName v) []++newTyVarWith :: Name -> TyVarInfo () -> TermM TyVar+newTyVarWith desc info = do+  i <- incCounter+  v <- newID $ mkTypeVarName desc i+  lvl <- curLevel+  modify $ \s -> s {termTyVars = M.insert v (lvl, info) $ termTyVars s}+  pure v++newTyVar :: (Located loc) => loc -> Liftedness -> Name -> TermM TyVar+newTyVar loc l desc = newTyVarWith desc $ TyVarFree (locOf loc) l++newType :: (Located loc) => loc -> Liftedness -> Name -> u -> TermM (TypeBase dim u)+newType loc l desc u = tyVarType u <$> newTyVar loc l desc++-- | New type that must be allowed as an array element.+newElemType :: (Located loc) => loc -> Name -> u -> TermM (TypeBase dim u)+newElemType loc desc u = tyVarType u <$> newTyVar loc Unlifted desc++newTypeWithField :: SrcLoc -> Name -> Name -> Type -> TermM Type+newTypeWithField loc desc k t =+  tyVarType NoUniqueness+    <$> newTyVarWith desc (TyVarRecord (locOf loc) $ M.singleton k t)++newTypeWithConstr :: SrcLoc -> Name -> u -> Name -> [TypeBase () u] -> TermM (TypeBase d u)+newTypeWithConstr loc desc u k ts =+  tyVarType u <$> newTyVarWith desc (TyVarSum (locOf loc) $ M.singleton k ts')+  where+    ts' = map (`setUniqueness` NoUniqueness) ts++newTypeOverloaded :: SrcLoc -> Name -> [PrimType] -> TermM (TypeBase d NoUniqueness)+newTypeOverloaded loc name pts =+  tyVarType NoUniqueness <$> newTyVarWith name (TyVarPrim (locOf loc) pts)++newArtificial :: u -> TypeBase () u -> TermM (TypeBase Size u)+newArtificial u t = do+  v <- newID "artificial"+  let t' = tyVarType u v+  modify $ \s -> s {termArtificial = M.insert v (second (const NoUniqueness) t) $ termArtificial s}+  pure t'++-- The AST requires annotations to be StructTypes, but the type+-- checker works with Types. This creates artificial type "variables"+-- that allow us to connect the AST annotations with the actual+-- inferred types. The artificial variables should never occur in+-- constraints - they can be substituted away with asType.+--+-- Equal components (with fully known shapes) of the same annotation+-- are given the same artificial variable, so that the sized type+-- checker knows that they have the same sizes.+asStructType :: TypeBase () u -> TermM (TypeBase Size u)+asStructType t = evalStateT (onType t) mempty+  where+    onType ::+      TypeBase () u' ->+      StateT (M.Map (TypeBase () NoUniqueness) TyVar) TermM (TypeBase Size u')+    onType (Scalar (Prim pt)) = pure $ Scalar $ Prim pt+    onType (Scalar (TypeVar u v [])) = pure $ Scalar $ TypeVar u v []+    onType (Scalar (Arrow u pname d t1 (RetType ext t2))) = do+      t1' <- onType t1+      t2' <- onType t2+      pure $ Scalar $ Arrow u pname d t1' $ RetType ext t2'+    onType (Scalar (Record fs)) =+      Scalar . Record <$> traverse onType fs+    onType (Scalar (Sum cs)) =+      Scalar . Sum <$> traverse (mapM onType) cs+    onType t'@(Scalar (TypeVar u _ _)) = artificial u t'+    onType t'@(Array u _ _) = artificial u t'++    artificial u t'+      | anonymousShape t' = lift $ newArtificial u t'+      | otherwise = do+          let key = second (const NoUniqueness) t'+          seen <- get+          case M.lookup key seen of+            Just v -> pure $ tyVarType u v+            Nothing -> do+              v <- lift $ newID "artificial"+              lift $ modify $ \s ->+                s {termArtificial = M.insert v key $ termArtificial s}+              modify $ M.insert key v+              pure $ tyVarType u v++    anonymousShape = elem () . bifoldMap (: []) (const mempty)++asType :: (Monoid u) => TypeBase Size u -> TermM (TypeBase () u)+asType t = do+  artificial <- gets termArtificial+  pure $ substTyVars (`M.lookup` artificial) (toType t)++expType :: Exp -> TermM Type+expType = asType . typeOf -- NOTE: Only place you should use typeOf.++addCt :: CtTy () -> TermM ()+addCt ct = modify $ \s -> s {termConstraints = ct : termConstraints s}++ctEq :: Reason (CtType ()) -> TypeBase () u1 -> TypeBase () u2 -> TermM ()+ctEq reason t1 t2 =+  -- As a minor optimisation, do not add constraint if the types are+  -- equal.+  unless (t1' == t2') $ addCt $ CtEq reason t1' t2'+  where+    t1' = t1 `setUniqueness` NoUniqueness+    t2' = t2 `setUniqueness` NoUniqueness++localScope :: (TermScope () -> TermScope ()) -> TermM a -> TermM a+localScope f = local $ \tenv -> tenv {termScope = f $ termScope tenv}++withEnv :: TermEnv -> Env -> TermEnv+withEnv tenv env = tenv {termScope = termScope tenv <> envToTermScope env}++lookupQualNameEnv :: QualName VName -> TermM (TermScope ())+lookupQualNameEnv qn = asks $ \tenv -> Scope.lookupQualNameEnv toType (termScope tenv) qn++instance MonadError TypeError TermM where+  throwError e = TermM $ do+    ws <- gets termWarnings+    throwError (ws, e)++  catchError (TermM m) f =+    TermM $ m `catchError` f'+    where+      f' (_, e) = let TermM m' = f e in m'++instance MonadTypeChecker TermM where+  warnings ws = modify $ \s -> s {termWarnings = termWarnings s <> ws}++  warn loc problem = warnings $ singleWarning (locOf loc) problem++  newName v = do+    s <- get+    let (v', src') = FreshNames.newName (termNameSource s) v+    put $ s {termNameSource = src'}+    pure v'++  newID s = newName $ VName s 0++  newTypeName name = do+    i <- incCounter+    newID $ mkTypeVarName name i++  bindVal v (TypeM.BoundV tps t) m = do+    t' <- asType t+    let f scope = scope {scopeVtable = M.insert v (BoundV tps t') $ scopeVtable scope}+    localScope f m++  lookupType qn = do+    outer_env <- asks termOuterEnv+    scope <- lookupQualNameEnv qn+    case M.lookup (qualLeaf qn) $ scopeTypeTable scope of+      Nothing -> error $ "lookupType: " <> show qn+      Just (TypeAbbr l ps (RetType dims def)) ->+        pure+          ( ps,+            RetType dims $ qualifyTypeVars outer_env (map typeParamName ps) (qualQuals qn) def,+            l+          )++  typeError loc notes s =+    throwError $ TypeError (locOf loc) notes s++--- All the general machinery goes above.++arrayOfRank :: Int -> Type -> Type+arrayOfRank n = arrayOf $ Shape $ replicate n ()++require :: T.Text -> [PrimType] -> Exp -> TermM Exp+require _why [pt] e = do+  e_t <- expType e+  ctEq (Reason (locOf e)) (Scalar $ Prim pt) e_t+  pure e+require _why pts e = do+  t :: Type <- newTypeOverloaded (srclocOf e) "t" pts+  e_t <- expType e+  ctEq (Reason (locOf e)) t e_t+  pure e++-- | Instantiate a type scheme with fresh type variables for its type+-- parameters. Returns the names of the fresh type variables, the+-- instance list, and the instantiated type.+instTypeScheme ::+  QualName VName ->+  SrcLoc ->+  [TypeParam] ->+  Type ->+  TermM ([VName], Type)+instTypeScheme qn loc tparams t = do+  (names, substs) <- fmap (unzip . catMaybes) $+    forM tparams $ \tparam ->+      case tparam of+        TypeParamType l v _ -> do+          v' <- newTyVar loc l $ nameFromText $ T.takeWhile isAscii $ nameToText $ baseName v+          modify $+            \s -> s {termTyInsts = TyInst (locOf loc) qn l v' : termTyInsts s}+          pure $ Just (v, (typeParamName tparam, tyVarType NoUniqueness v'))+        TypeParamDim {} ->+          pure Nothing+  let t' = substTyVars (`lookup` substs) t+  pure (names, t')++lookupMod :: QualName VName -> TermM Mod+lookupMod qn@(QualName _ name) = do+  scope <- lookupQualNameEnv qn+  case M.lookup name $ scopeModTable scope of+    Nothing -> error $ "lookupMod: " <> show qn+    Just m -> pure m++lookupVar :: SrcLoc -> QualName VName -> TermM Type+lookupVar loc qn@(QualName qs name) = do+  scope <- lookupQualNameEnv qn+  outer_env <- asks termOuterEnv+  -- Top-level value bindings are not in the term scope (see+  -- 'envToTermScopeNoVals'); look them up on demand in the outer+  -- environment.+  case M.lookup name (scopeVtable scope) `mplus` Scope.lookupOuterVal toType outer_env name of+    Nothing ->+      error $ "lookupVar: " <> show qn+    Just (BoundV tparams t) -> do+      if null tparams && null qs+        then pure t+        else do+          (tnames, t') <- instTypeScheme qn loc tparams t+          -- Qualify abstract types, so that e.g. mismatch errors+          -- mention them by how they were accessed. The sizes need+          -- no qualification, as they are not even present here.+          pure $ qualifyTypeVarsWith (\_ _ d -> d) outer_env tnames qs t'+    Just EqualityF -> do+      argtype <- tyVarType Observe <$> newTyVarWith "t" (TyVarFree (locOf loc) Unlifted)+      pure $ foldFunType [argtype, argtype] $ RetType [] $ Scalar $ Prim Bool+    Just (OverloadedF ts pts rt) -> do+      argtype <- newTypeOverloaded loc "t" ts+      let (pts', rt') = instOverloaded argtype pts rt+      pure $ foldFunType (map (second $ const Observe) pts') $ RetType [] $ second (const Nonunique) rt'+    -- The unsized checker binds recursive functions directly (see+    -- 'checkRecursive'), so it never produces a 'RecursiveV'.+    Just RecursiveV ->+      error $ "lookupVar: unexpected RecursiveV for " <> show qn+  where+    instOverloaded argtype pts rt =+      ( map (maybe argtype (Scalar . Prim)) pts,+        maybe argtype (Scalar . Prim) rt+      )++bind ::+  [Ident StructType] ->+  TermM a ->+  TermM a+bind idents m = do+  let names = map identName idents+  ts <- mapM (asType . unInfo . identType) idents+  localScope (`bindVars` zip names ts) m+  where+    bindVars = foldl bindVar++    bindVar scope (name, t) =+      scope+        { scopeVtable = M.insert name (BoundV [] t) $ scopeVtable scope+        }++-- All this complexity is just so we can handle un-suffixed numeric+-- literals in patterns.+patLitMkType :: PatLit -> SrcLoc -> TermM ParamType+patLitMkType (PatLitInt _) loc =+  toParam Observe <$> newTypeOverloaded loc "t" anyNumberType+patLitMkType (PatLitFloat _) loc =+  toParam Observe <$> newTypeOverloaded loc "t" anyFloatType+patLitMkType (PatLitPrim v) _ =+  pure $ Scalar $ Prim $ primValueType v++checkSizeExp' :: ExpBase NoInfo VName -> TermM Exp+checkSizeExp' e = do+  e' <- checkExp e+  e_t <- expType e'+  ctEq (Reason (locOf e)) e_t (Scalar (Prim (Signed Int64)))+  pure e'++checkPat' ::+  PatBase NoInfo VName ParamType ->+  Inferred ParamType ->+  TermM (Pat ParamType)+checkPat' (PatParens p loc) t =+  PatParens <$> checkPat' p t <*> pure loc+checkPat' (PatAttr attr p loc) t =+  PatAttr <$> checkAttr attr <*> checkPat' p t <*> pure loc+checkPat' (Id name NoInfo loc) (Ascribed t) =+  pure $ Id name (Info t) loc+checkPat' (Id name NoInfo loc) NoneInferred = do+  t <- newType loc Lifted "t" Observe+  pure $ Id name (Info t) loc+checkPat' (Wildcard _ loc) (Ascribed t) = do+  pure $ Wildcard (Info t) loc+checkPat' (Wildcard NoInfo loc) NoneInferred = do+  t <- newType loc Lifted "t" Observe+  pure $ Wildcard (Info t) loc+checkPat' p@(TuplePat ps loc) (Ascribed t)+  | Just ts <- isTupleRecord t,+    length ts == length ps =+      TuplePat+        <$> zipWithM checkPat' ps (map Ascribed ts)+        <*> pure loc+  | otherwise =+      typeError loc mempty $+        "Pattern"+          </> indent 2 (pretty p)+          </> "cannot match ascribed type"+          </> indent 2 (pretty t)+checkPat' (TuplePat ps loc) NoneInferred =+  TuplePat <$> mapM (`checkPat'` NoneInferred) ps <*> pure loc+checkPat' p@(RecordPat p_fs loc) _+  | Just (L floc f, _) <- L.find (("_" `T.isPrefixOf`) . nameToText . unLoc . fst) p_fs =+      typeError floc mempty $+        "Underscore-prefixed fields are not allowed."+          </> "Did you mean"+          <> dquotes (pretty (T.drop 1 (nameToText f)) <> "=_")+          <> "?"+  | length (nubOrd (map fst p_fs)) /= length (map fst p_fs) =+      typeError loc mempty $+        "Duplicate fields in record pattern" <+> pretty p <> "."+checkPat' p@(RecordPat p_fs loc) (Ascribed t)+  | Scalar (Record t_fs) <- t,+    p_fs' <- L.sortBy (comparing fst) p_fs,+    t_fs' <- L.sortBy (comparing fst) (M.toList t_fs),+    map fst t_fs' == map (unLoc . fst) p_fs' =+      RecordPat <$> zipWithM check p_fs' t_fs' <*> pure loc+  | otherwise = do+      typeError loc mempty $+        "Pattern"+          </> indent 2 (pretty p)+          </> "cannot match ascribed type"+          </> indent 2 (pretty t)+  where+    check (L f_loc f, p_f) (_, t_f) =+      (L f_loc f,) <$> checkPat' p_f (Ascribed t_f)+checkPat' (RecordPat fs loc) NoneInferred =+  RecordPat . M.toList+    <$> traverse (`checkPat'` NoneInferred) (M.fromList fs)+    <*> pure loc+checkPat' (PatAscription p t loc) maybe_outer_t = do+  (t', _, RetType _ st, _) <- checkTypeExp checkSizeExp' t++  let st' = resToParam st++  case maybe_outer_t of+    Ascribed outer_t -> do+      unless (toType st' == toType outer_t) $+        typeError loc mempty $+          "Ascribed type"+            </> indent 2 (pretty st)+            </> "cannot match outer ascribed type"+            </> indent 2 (pretty outer_t)+      PatAscription+        <$> checkPat' p (Ascribed st')+        <*> pure t'+        <*> pure loc+    NoneInferred ->+      PatAscription+        <$> checkPat' p (Ascribed st')+        <*> pure t'+        <*> pure loc+checkPat' (PatLit l NoInfo loc) (Ascribed t) = do+  t' <- patLitMkType l loc+  ctEq (Reason (locOf loc)) (toType t') (toType t)+  pure $ PatLit l (Info t') loc+checkPat' (PatLit l NoInfo loc) NoneInferred = do+  t' <- patLitMkType l loc+  pure $ PatLit l (Info t') loc+checkPat' (PatConstr n NoInfo ps loc) (Ascribed (Scalar (Sum cs)))+  | Just ts <- M.lookup n cs,+    length ps == length ts = do+      ps' <- zipWithM checkPat' ps $ map Ascribed ts+      pure $ PatConstr n (Info (Scalar (Sum cs))) ps' loc+checkPat' p@(PatConstr {}) (Ascribed t) =+  typeError (locOf p) mempty $+    "Pattern"+      </> indent 2 (pretty p)+      </> "cannot match ascribed type"+      </> indent 2 (pretty t)+checkPat' (PatConstr n NoInfo ps loc) NoneInferred = do+  ps' <- mapM (`checkPat'` NoneInferred) ps+  t <- newTypeWithConstr loc "t" Observe n =<< mapM (asType . patternType) ps'+  pure $ PatConstr n (Info $ toParam Observe t) ps' loc++checkPat ::+  PatBase NoInfo VName (TypeBase Size u) ->+  (Pat ParamType -> TermM a) ->+  TermM a+checkPat p m =+  m =<< checkPat' (fmap (toParam Observe) p) NoneInferred++-- | Bind @let@-bound sizes. This is usually followed by 'bindLetPat'+-- immediately afterwards.+bindSizes :: [SizeBinder VName] -> TermM a -> TermM a+bindSizes [] m = m -- Minor optimisation.+bindSizes sizes m = bind (map sizeWithType sizes) m+  where+    sizeWithType size =+      Ident (sizeName size) (Info (Scalar (Prim (Signed Int64)))) (srclocOf size)++bindLetPat ::+  PatBase NoInfo VName (TypeBase Size u) ->+  Type ->+  (Pat ParamType -> TermM a) ->+  TermM a+bindLetPat p t m = do+  checkPat p $ \p' -> do+    pt <- asType $ patternType p'+    ctEq (ReasonPatMatch (locOf p) (fmap toStruct p) t) pt t+    bind (patIdents (fmap toStruct p')) $ m p'++bindTypes ::+  [(VName, TypeBinding)] ->+  TermM a ->+  TermM a+bindTypes tbinds = localScope extend+  where+    extend scope =+      scope+        { scopeTypeTable = M.fromList tbinds <> scopeTypeTable scope+        }++bindTypeParams :: [TypeParam] -> TermM a -> TermM a+bindTypeParams tparams m =+  bind idents . bindTypes types $ do+    lvl <- curLevel+    modify $ \s ->+      s+        { termTyParams =+            termTyParams s+              <> M.fromList (mapMaybe (typeParam lvl) tparams)+        }+    m+  where+    idents = mapMaybe typeParamIdent tparams+    types = mapMaybe typeParamType tparams+    typeParamType (TypeParamType l v _) =+      Just (v, TypeAbbr l [] $ RetType [] $ Scalar (TypeVar mempty (qualName v) []))+    typeParamType TypeParamDim {} = Nothing+    typeParam lvl (TypeParamType l v loc) = Just (v, (lvl, l, locOf loc))+    typeParam _ _ = Nothing++bindParams ::+  [TypeParam] ->+  [PatBase NoInfo VName ParamType] ->+  ([Pat ParamType] -> TermM a) ->+  TermM a+bindParams tps orig_ps m = bindTypeParams tps $ do+  let descend ps' (p : ps) =+        checkPat p $ \p' ->+          bind (patIdents $ fmap toStruct p') $ incLevel $ descend (p' : ps') ps+      descend ps' [] = m $ reverse ps'++  incLevel $ descend [] orig_ps++checkApplyOne ::+  SrcLoc ->+  (Maybe (QualName VName), Int) ->+  Type ->+  (Maybe Exp, Type) ->+  TermM Type+checkApplyOne loc fname ftype (arg, argtype) = do+  (a, b) <- split ftype+  let reason = case arg of+        Just arg' -> ReasonApply (locOf arg) fname arg' a argtype+        Nothing -> Reason (locOf loc)+  ctEq reason argtype a+  pure b+  where+    split (Scalar (Arrow _ _ _ a (RetType _ b))) =+      pure (a, b `setUniqueness` NoUniqueness)+    split (Array _u s t) = do+      (a, b) <- split $ Scalar t+      pure (arrayOf s a, arrayOf s b)+    split ftype' = do+      a <- newType loc Lifted "arg" NoUniqueness+      b <- newType loc Lifted "res" Nonunique+      let reason = case arg of+            Just arg' -> ReasonApplySplit (locOf loc) fname arg' ftype'+            Nothing -> Reason $ locOf loc+      ctEq reason ftype' $ Scalar $ Arrow NoUniqueness Unnamed Observe a $ RetType [] b+      pure (a, b `setUniqueness` NoUniqueness)++checkApply ::+  SrcLoc ->+  Maybe (QualName VName) ->+  Type ->+  NE.NonEmpty (Maybe Exp, Type) ->+  TermM Type+checkApply loc fname ftype args = do+  (_, rt) <- foldM onArg (0, ftype) args+  pure rt+  where+    onArg (i, f_t) arg = do+      rt <- checkApplyOne loc (fname, i) f_t arg+      pure (i + 1, rt)++checkSlice :: SliceBase NoInfo VName -> TermM [DimIndex]+checkSlice = mapM checkDimIndex+  where+    checkDimIndex (DimFix i) =+      DimFix <$> (require "use as index" anySignedType =<< checkExp i)+    checkDimIndex (DimSlice i j s) =+      DimSlice <$> traverse check i <*> traverse check j <*> traverse check s++    check = require "use in slice" [Signed Int64] <=< checkExp++isSlice :: DimIndexBase f vn -> Bool+isSlice DimSlice {} = True+isSlice DimFix {} = False++checkCase ::+  Type ->+  CaseBase NoInfo VName ->+  TermM (CaseBase Info VName, Type)+checkCase mt (CasePat p e loc) =+  bindLetPat p mt $ \p' -> do+    e' <- checkExp e+    e_t <- expType e'+    pure (CasePat (fmap toStruct p') e' loc, e_t)++checkCases ::+  Type ->+  NE.NonEmpty (CaseBase NoInfo VName) ->+  TermM (NE.NonEmpty (CaseBase Info VName), Type)+checkCases mt rest_cs = do+  let (c, rest_cs') = NE.uncons rest_cs+  (c', c_t) <- checkCase mt c+  case rest_cs' of+    Nothing ->+      pure (NE.singleton c', c_t)+    Just cs -> do+      (cs', cs_t) <- checkCases mt cs+      ctEq (ReasonBranches (locOf c) c_t cs_t) c_t cs_t+      pure (NE.cons c' cs', c_t)++checkRetDecl ::+  Exp ->+  Maybe (TypeExp (ExpBase NoInfo VName) VName) ->+  TermM (Type, Maybe (TypeExp Exp VName))+checkRetDecl body Nothing = (,Nothing) <$> expType body+checkRetDecl body (Just te) = do+  (te', _, RetType _ st, _) <- checkTypeExp checkSizeExp' te+  body_t <- expType body+  st' <- toStruct <$> asType st+  ctEq (ReasonRetType (locOf body) st' body_t) st' body_t+  pure (st', Just te')++-- Add constraints saying that the first type has a (potentially nested) part+-- containing the second type.+--+-- FIXME: the locations here are very bad.+mustHaveSteps ::+  (Pretty a, Located a) =>+  a ->+  Type ->+  [UpdateStep Info VName] ->+  Type ->+  TermM ()+mustHaveSteps src t [] ve_t =+  -- This case is probably never reached.+  ctEq (Reason (locOf src)) t ve_t+mustHaveSteps src t [UpdateStepField f] ve_t = do+  rt :: Type <- newTypeWithField (srclocOf src) "ft" f ve_t+  ctEq (Reason (locOf src)) t rt+mustHaveSteps src t (UpdateStepField f : steps) ve_t = do+  ft <- newType (locOf src) Lifted "ft" NoUniqueness+  rt :: Type <- newTypeWithField (srclocOf src) "ft" f ft+  ctEq (Reason (locOf src)) t rt+  mustHaveSteps src ft steps ve_t+mustHaveSteps src t [UpdateStepSlice slice] ve_t = do+  let num_slices = length $ filter isSlice slice+  update_elem_t <- newElemType (locOf src) "update_elem" NoUniqueness+  ctEq (Reason (locOf src)) t $ arrayOfRank (length slice) update_elem_t+  ctEq (Reason (locOf src)) ve_t $ arrayOfRank num_slices update_elem_t+mustHaveSteps src t (UpdateStepSlice slice : steps) ve_t = do+  let num_slices = length $ filter isSlice slice+  index_tv <- newTyVar (locOf src) Unlifted "index"+  index_elem_t <- newElemType (locOf src) "index_elem" NoUniqueness+  ctEq (Reason (locOf src)) (tyVarType NoUniqueness index_tv) $ arrayOfRank num_slices index_elem_t+  ctEq (Reason (locOf src)) t $ arrayOfRank (length slice) index_elem_t+  mustHaveSteps src (arrayOfRank num_slices index_elem_t) steps ve_t++checkStep :: UpdateStep NoInfo VName -> TermM (UpdateStep Info VName)+checkStep (UpdateStepField f) = pure $ UpdateStepField f+checkStep (UpdateStepSlice slice) = UpdateStepSlice <$> checkSlice slice++checkExp :: ExpBase NoInfo VName -> TermM (ExpBase Info VName)+--+checkExp (Var qn _ loc) = do+  t <- asStructType =<< lookupVar loc qn+  pure $ Var qn (Info t) loc+checkExp (OpSection op _ loc) = do+  ftype <- asStructType =<< lookupVar loc op+  pure $ OpSection op (Info ftype) loc+checkExp (Negate arg loc) = do+  arg' <- require "numeric negation" anyNumberType =<< checkExp arg+  pure $ Negate arg' loc+checkExp (Not arg loc) = do+  arg' <- require "logical negation" (Bool : anyIntType) =<< checkExp arg+  pure $ Not arg' loc+checkExp (Hole NoInfo loc) =+  Hole <$> (Info <$> newType loc Lifted "hole" NoUniqueness) <*> pure loc+checkExp (Parens e loc) =+  Parens <$> checkExp e <*> pure loc+checkExp (TupLit es loc) =+  TupLit <$> mapM checkExp es <*> pure loc+checkExp (QualParens (modname, modnameloc) e loc) = do+  mod <- lookupMod modname+  case mod of+    ModEnv env -> local (`withEnv` env) $ do+      e' <- checkExp e+      pure $ QualParens (modname, modnameloc) e' loc+    ModFun {} ->+      typeError loc mempty . withIndexLink "module-is-parametric" $+        "Module" <+> pretty modname <+> " is a parametric module."+--+checkExp (IntLit x NoInfo loc) = do+  t <- newTypeOverloaded loc "num" anyNumberType+  pure $ IntLit x (Info t) loc+checkExp (FloatLit x NoInfo loc) = do+  t <- newTypeOverloaded loc "float" anyFloatType+  pure $ FloatLit x (Info t) loc+checkExp (Literal v loc) =+  pure $ Literal v loc+checkExp (StringLit vs loc) =+  pure $ StringLit vs loc+-- No need to type check this, as these are only produced by the+-- parser if the elements are monomorphic and all match.+checkExp (ArrayVal vs t loc) =+  pure $ ArrayVal vs t loc+checkExp (ArrayLit es _ loc) = do+  -- TODO: this will produce an enormous number of constraints and+  -- type variables for pathologically large arrays with+  -- type-unsuffixed integers. Add some special case that handles that+  -- more efficiently.+  et <- newElemType loc "et" NoUniqueness+  es' <- forM es $ \e -> do+    e' <- checkExp e+    e_t <- expType e'+    et' <- asType et+    ctEq (Reason (locOf loc)) e_t et'+    pure e'+  let arr_t = arrayOf (Shape [sizeFromInteger (L.genericLength es) loc]) et+  pure $ ArrayLit es' (Info arr_t) loc+checkExp (RecordLit fs loc) =+  RecordLit <$> evalStateT (mapM checkField fs) mempty <*> pure loc+  where+    checkField (RecordFieldExplicit f e rloc) = do+      errIfAlreadySet (unLoc f) rloc+      modify $ M.insert (unLoc f) rloc+      RecordFieldExplicit f <$> lift (checkExp e) <*> pure rloc+    checkField (RecordFieldImplicit name NoInfo rloc) = do+      errIfAlreadySet (baseName (unLoc name)) rloc+      t <- lift $ asStructType =<< lookupVar rloc (qualName (unLoc name))+      modify $ M.insert (baseName (unLoc name)) rloc+      pure $ RecordFieldImplicit name (Info t) rloc++    errIfAlreadySet f rloc = do+      maybe_sloc <- gets $ M.lookup f+      case maybe_sloc of+        Just sloc ->+          lift . typeError rloc mempty $+            "Field"+              <+> dquotes (pretty f)+              <+> "previously defined at"+              <+> pretty (locStrRel rloc sloc)+              <> "."+        Nothing -> pure ()++--+checkExp (Attr info e loc) =+  Attr <$> checkAttr info <*> checkExp e <*> pure loc+checkExp (Assert e1 e2 NoInfo loc) = do+  e1' <- require "being asserted" [Bool] =<< checkExp e1+  e2' <- checkExp e2+  pure $ Assert e1' e2' (Info (prettyText e1)) loc+--+checkExp (Constr name es NoInfo loc) = do+  es' <- mapM checkExp es+  es_ts <- mapM expType es'+  t <- newTypeWithConstr loc "t" NoUniqueness name es_ts+  pure $ Constr name es' (Info t) loc+--+checkExp (AppExp (Apply fe args loc) NoInfo) = do+  fe' <- checkExp fe+  (args', apply_args) <-+    fmap NE.unzip . forM args $ \(_, arg) -> do+      arg' <- checkExp arg+      arg_t <- expType arg'+      pure (arg', (Just arg', arg_t))+  fe_t <- expType fe'+  rt <- checkApply loc fname fe_t apply_args+  rt' <- asStructType rt+  let args'' = NE.map (\arg -> (Info Nothing, arg)) args'+  pure $ AppExp (Apply fe' args'' loc) $ Info (AppRes rt' [])+  where+    fname =+      case fe of+        Var v _ _ -> Just v+        _ -> Nothing+checkExp (AppExp (BinOp (op, oploc) NoInfo (e1, _) (e2, _) loc) NoInfo) = do+  ftype <- lookupVar oploc op+  e1' <- checkExp e1+  e1_t <- expType e1'+  e2' <- checkExp e2+  e2_t <- expType e2'++  rt <-+    checkApply+      loc+      (Just op)+      ftype+      ((Just e1', e1_t) NE.:| [(Just e2', e2_t)])+  rt' <- asStructType rt++  ftype' <- asStructType ftype+  pure $+    AppExp+      (BinOp (op, oploc) (Info ftype') (e1', Info Nothing) (e2', Info Nothing) loc)+      (Info (AppRes rt' []))+--+checkExp (OpSectionLeft op _ e _ _ loc) = do+  optype <- lookupVar loc op+  e' <- checkExp e+  e_t <- expType e'+  t2 <- newType loc Lifted "t" NoUniqueness+  t2' <- asStructType t2+  rt <-+    checkApply+      loc+      (Just op)+      optype+      ((Just e', e_t) NE.:| [(Nothing, t2)])+  rt' <- asStructType rt++  t1 <- asStructType e_t+  optype' <- asStructType optype+  pure $+    OpSectionLeft+      op+      (Info optype')+      e'+      ( Info (Unnamed, toParam Observe t1, Nothing),+        Info (Unnamed, toParam Observe t2')+      )+      (Info (RetType [] (rt' `setUniqueness` Nonunique)), Info [])+      loc+checkExp (OpSectionRight op _ e _ NoInfo loc) = do+  optype <- lookupVar loc op+  e' <- checkExp e+  e_t <- expType e'+  t1 <- newType loc Lifted "t" NoUniqueness+  t1' <- asStructType t1+  rt <-+    checkApply+      loc+      (Just op)+      optype+      ((Nothing, t1) NE.:| [(Just e', e_t)])+  rt' <- asStructType rt+  t2 <- asStructType e_t++  optype' <- asStructType optype+  pure $+    OpSectionRight+      op+      (Info optype')+      e'+      -- Dummy types.+      ( Info (Unnamed, toParam Observe t1'),+        Info (Unnamed, toParam Observe t2, Nothing)+      )+      (Info $ RetType [] (rt' `setUniqueness` Nonunique))+      loc+--+checkExp e@(UpdateSection steps NoInfo loc) = do+  steps' <- mapM checkStep steps+  -- Lifted, as a pure field projection works on records with+  -- function-typed fields. Any slice steps will constrain the+  -- relevant parts to be arrays (of unlifted elements) anyway.+  src_t <- newType loc Lifted "update" NoUniqueness+  ve_t <- newType loc Lifted "update_elem" NoUniqueness+  mustHaveSteps e src_t steps' ve_t+  ft <-+    asStructType $+      Scalar $+        Arrow mempty Unnamed Observe src_t $+          second (const Nonunique) (RetType [] ve_t)+  pure $ UpdateSection steps' (Info ft) loc++--+checkExp (Lambda params body retdecl NoInfo loc) = do+  bindParams [] params $ \params' -> do+    body' <- checkExp body++    (body_t, retdecl') <- checkRetDecl body' retdecl+    body_t' <- asStructType body_t+    let ret = RetType [] $ toRes Nonunique body_t'+    pure $ Lambda params' body' retdecl' (Info ret) loc+--+checkExp (AppExp (LetPat sizes pat e body loc) _) = do+  e' <- checkExp e+  e_t <- expType e'++  bindSizes sizes . incLevel . bindLetPat pat e_t $ \pat' -> do+    body' <- incLevel $ checkExp body+    body_t <- expType body'++    body_t' <- asStructType body_t+    pure $+      AppExp+        (LetPat sizes (fmap toStruct pat') e' body' loc)+        (Info $ AppRes body_t' [])+--+checkExp (AppExp (LetFun name (tparams, params, retdecl, NoInfo, e) body loc) _) = do+  (tparams', params', retdecl', rettype, e') <-+    bindParams tparams params $ \params' -> do+      e' <- checkExp e+      (e_t, retdecl') <- checkRetDecl e' retdecl+      pure (tparams, params', retdecl', fmap (const Nonunique) e_t, e')++  params'' <- mapM (traverse asType) params'++  let entry = BoundV tparams' $ funType params'' $ RetType [] rettype+      bindF scope =+        scope+          { scopeVtable = M.insert (fst name) entry $ scopeVtable scope+          }+  body' <- localScope bindF $ checkExp body+  body_t <- expType body'++  body_t' <- asStructType body_t+  rettype' <- asStructType rettype+  pure $+    AppExp+      ( LetFun+          name+          (tparams', params', retdecl', Info (RetType [] rettype'), e')+          body'+          loc+      )+      (Info $ AppRes body_t' [])+--+checkExp (AppExp (Range start maybe_step end loc) _) = do+  start' <- require "use in range expression" anyIntType =<< checkExp start+  let check e = do+        e' <- checkExp e+        start_t <- expType start'+        e_t <- expType e'+        ctEq (Reason (locOf e')) start_t e_t+        pure e'+  maybe_step' <- traverse check maybe_step+  end' <- traverse check end+  range_t <- newElemType loc "range" NoUniqueness+  range_t' <- asType range_t+  start_t <- expType start'+  ctEq (Reason (locOf start')) range_t' (arrayOfRank 1 start_t)+  pure $ AppExp (Range start' maybe_step' end' loc) $ Info $ AppRes range_t []+--+checkExp (Project k e NoInfo loc) = do+  e' <- checkExp e+  kt <- newType loc Lifted "kt" NoUniqueness+  t <- newTypeWithField loc "t" k kt+  e_t <- expType e'+  ctEq (Reason (locOf e')) e_t t+  kt' <- asStructType kt+  pure $ Project k e' (Info kt') loc+--+checkExp (Update src steps ve NoInfo loc) = do+  src' <- checkExp src+  src_t <- expType src'+  src_t' <- asStructType src_t+  ve' <- checkExp ve+  ve_t <- expType ve'+  steps' <- mapM checkStep steps+  mustHaveSteps src' src_t steps' ve_t+  pure $ Update src' steps' ve' (Info src_t') loc++--+checkExp (AppExp (Index e slice loc) _) = do+  e' <- checkExp e+  e_t <- expType e'+  slice' <- checkSlice slice+  index_tv <- newTyVar loc Unlifted "index"+  index_elem_t <- newElemType loc "index_elem" NoUniqueness+  let num_slices = length $ filter isSlice slice+  ctEq (Reason (locOf loc)) (tyVarType NoUniqueness index_tv) $ arrayOfRank num_slices index_elem_t+  ctEq (Reason (locOf e')) e_t $ arrayOfRank (length slice) index_elem_t+  pure $ AppExp (Index e' slice' loc) (Info $ AppRes (tyVarType NoUniqueness index_tv) [])+--+checkExp (AppExp (LetWith dest src steps ve body loc) _) = do+  src_t <- lookupVar (srclocOf src) $ qualName $ identName src+  src_t' <- asStructType src_t+  let src' = src {identType = Info src_t'}+      dest' = dest {identType = Info src_t'}+  steps' <- mapM checkStep steps+  ve' <- checkExp ve+  ve_t <- expType ve'+  mustHaveSteps src' src_t steps' ve_t+  bind [dest'] $ do+    body' <- checkExp body+    body_t <- expType body'+    body_t' <- asStructType body_t+    pure $ AppExp (LetWith dest' src' steps' ve' body' loc) (Info $ AppRes body_t' [])+--+checkExp (AppExp (If e1 e2 e3 loc) _) = do+  e1' <- checkExp e1+  e1_t <- expType e1'+  e2' <- checkExp e2+  e2_t <- expType e2'+  e3' <- checkExp e3+  e3_t <- expType e3'+  if_t <- newType loc SizeLifted "if_t" NoUniqueness++  ctEq (Reason (locOf e1')) e1_t (Scalar (Prim Bool))+  ctEq (ReasonBranches (locOf loc) e2_t e3_t) e2_t if_t+  ctEq (ReasonBranches (locOf loc) e2_t e3_t) e3_t if_t++  if_t' <- asStructType if_t+  pure $ AppExp (If e1' e2' e3' loc) (Info $ AppRes if_t' [])+--+checkExp (AppExp (Match e cs loc) _) = do+  e' <- checkExp e+  e_t <- expType e'+  (cs', t) <- checkCases e_t cs++  match_t <- newType loc SizeLifted "match_t" NoUniqueness+  ctEq (Reason (locOf loc)) match_t t++  match_t' <- asStructType match_t+  pure $ AppExp (Match e' cs' loc) (Info $ AppRes match_t' [])+--+checkExp (AppExp (Loop _ pat arg form body loc) _) = do+  arg' <- checkExp $ case arg of+    LoopInitExplicit e -> e+    LoopInitImplicit _ ->+      -- Should have been filled out in Names+      error "Unspected LoopInitImplicit"+  arg_t <- expType arg'+  loop_t <- newType loc SizeLifted "loop_t" NoUniqueness+  ctEq (Reason (locOf loc)) arg_t loop_t+  bindLetPat pat arg_t $ \pat' -> do+    (form', body') <-+      case form of+        For (Ident i _ iloc) bound -> do+          bound' <- require "loop bound" anyIntType =<< checkExp bound+          bound_t <- expType bound'+          bound_t' <- asStructType bound_t+          let i' = Ident i (Info bound_t') iloc+          bind [i'] $ do+            body' <- checkExp body+            pure (For i' bound', body')+        While cond -> do+          cond' <- checkExp cond+          body' <- checkExp body+          pure (While cond', body')+        ForIn elemp arr -> do+          arr' <- checkExp arr+          elem_t <- newElemType elemp "elem" NoUniqueness+          arr_t <- expType arr'+          elem_t' <- asType elem_t+          ctEq (Reason (locOf arr')) arr_t $ arrayOfRank 1 elem_t'+          bindLetPat elemp elem_t' $ \elemp' -> do+            body' <- checkExp body+            pure (ForIn (toStruct <$> elemp') arr', body')+    body_t <- expType body'+    ctEq (Reason (locOf loc)) arg_t body_t+    pure $+      AppExp+        (Loop [] pat' (LoopInitExplicit arg') form' body' loc)+        (Info (AppRes (patternStructType pat') []))+--+checkExp (Ascript e te loc) = do+  e' <- checkExp e+  (te', _, RetType _ st, _) <- checkTypeExp checkSizeExp' te+  e_t <- expType e'+  st' <- asType st+  ctEq (ReasonAscription (locOf e') (toStruct st') (toStruct e_t)) e_t st'+  pure $ Ascript e' te' loc+checkExp (Coerce e te NoInfo loc) = do+  e' <- checkExp e+  (te', _, RetType _ st, _) <- checkTypeExp checkSizeExp' te+  e_t <- expType e'+  st' <- asType st+  ctEq (Reason (locOf e')) e_t st'+  pure $ Coerce e' te' (Info (toStruct st)) loc++doDefault ::+  [VName] ->+  VName ->+  Either [PrimType] (TypeBase () NoUniqueness) ->+  TermM (TypeBase () NoUniqueness)+doDefault tyvars_at_toplevel v (Left pts)+  | [pt] <- pts =+      pure $ Scalar $ Prim pt+  | Signed Int32 `elem` pts = do+      when (v `elem` tyvars_at_toplevel) $+        warn usage "Defaulting ambiguous type to i32."+      pure $ Scalar $ Prim $ Signed Int32+  | FloatType Float64 `elem` pts = do+      when (v `elem` tyvars_at_toplevel) $+        warn usage "Defaulting ambiguous type to f64."+      pure $ Scalar $ Prim $ FloatType Float64+  | otherwise =+      typeError usage mempty . withIndexLink "ambiguous-type" $+        "Type is ambiguous (could be one of"+          <+> commasep (map pretty pts)+          <> ")."+            </> "Add a type annotation to disambiguate the type."+  where+    usage = mkUsage NoLoc "overload"+doDefault _ _ (Right t) = pure t++-- | Apply defaults on otherwise ambiguous types. This may result in+-- some type variables becoming known, so we have to perform+-- substitutions on the RHS of the substitutions afterwards.+doDefaults ::+  [VName] ->+  M.Map TyVar (Either [PrimType] (TypeBase () NoUniqueness)) ->+  TermM (M.Map TyVar (TypeBase () NoUniqueness))+doDefaults tyvars_at_toplevel substs = do+  substs' <- M.traverseWithKey (doDefault tyvars_at_toplevel) substs+  pure $ M.map (substTyVars (`M.lookup` substs')) substs'++generalise ::+  TypeBase () NoUniqueness ->+  [UnconTyVar] ->+  Solution ->+  ([TypeParam], [VName])+generalise fun_t unconstrained solution =+  -- Candidates for let-generalisation are those type variables that+  -- are used in fun_t.+  let visible = foldMap expandTyVars $ typeVars fun_t+      onTyVar (v, l)+        | v `S.member` visible = Left $ TypeParamType l v mempty+        | otherwise = Right v+   in partitionEithers $ map onTyVar unconstrained+  where+    expandTyVars v =+      case M.lookup v solution of+        Just (Right t) -> foldMap expandTyVars $ typeVars t+        _ -> S.singleton v++generaliseAndDefaults ::+  [UnconTyVar] ->+  Solution ->+  TypeBase () NoUniqueness ->+  TermM ([TypeParam], M.Map VName (TypeBase () NoUniqueness))+generaliseAndDefaults unconstrained solution t = do+  let (generalised, unconstrained') =+        generalise t unconstrained solution+      -- See #1552 for why we resolve unconstrained and un-generalised type+      -- variables to ().+      units = M.fromList (map (,Right (Scalar (Record mempty))) unconstrained')+  solution' <- doDefaults (S.toList $ typeVars t) (units <> solution)+  pure+    ( generalised,+      solution'+    )++-- | Verify that the recorded type parameter instantiations respect the+-- liftedness of the type parameters. The constraint solver merely propagates+-- liftedness constraints; this is where they are enforced for instantiations,+-- as only here do we know why the constraints exist. (Other liftedness rules+-- are enforced by 'localChecks' in Language.Futhark.TypeChecker.Terms.)+checkTyInstLiftedness :: Solution -> TermM ()+checkTyInstLiftedness solution = do+  typarams <- gets termTyParams+  tyset <- asks termTySet+  mapM_ (check typarams tyset) . reverse =<< gets termTyInsts+  where+    -- A Lifted type parameter permits any instantiation.+    check _ _ (TyInst _ _ Lifted _) = pure ()+    check typarams tyset (TyInst loc qn l v)+      | Just (Right t) <- M.lookup v solution = do+          unless (orderZero t) . typeError loc mempty $+            "Type"+              </> indent 2 (pretty t)+              </> "found to be functional."+              </> when_inst+          let bad = case l of+                Unlifted -> [Lifted, SizeLifted]+                _ -> [Lifted]+              -- The liftedness of a type variable is given by its+              -- binding if it is a type parameter, and by the type+              -- set if it is an abstract type.+              badVar qv =+                case M.lookup (qualLeaf qv) typarams of+                  Just (_, pl, ploc) -> do+                    guard $ pl `elem` bad+                    Just $+                      "Type parameter"+                        <+> dquotes (prettyName (qualLeaf qv))+                        <+> "bound at"+                        <+> pretty (locStr ploc)+                  Nothing -> do+                    al <- M.lookup qv tyset+                    guard $ al `elem` bad+                    Just $ "Type" <+> dquotes (pretty qv)+          case mapMaybe badVar $ typeQualVars t of+            what : _ ->+              typeError loc mempty $+                what+                  <+> case l of+                    Unlifted -> "is lifted and cannot be an array element."+                    _ -> "is lifted and may be a functional type."+                  </> when_inst+            [] -> pure ()+      | otherwise = pure ()+      where+        when_inst =+          "When instantiating type parameter of" <+> dquotes (pretty qn) <> "."++-- | Check a potentially recursive function body. The function is bound to a+-- fresh monomorphic type variable while its body is checked; that variable is+-- then constrained to the actual function type, and the constraint solver ties+-- the knot. A parameterless binding cannot be recursive (see 'resolveValBind'),+-- so it is checked with no self-reference in scope. See Note [Checking recursive+-- functions] in Language.Futhark.TypeChecker.Terms.+checkRecursive ::+  VName ->+  SrcLoc ->+  [Pat ParamType] ->+  ExpBase NoInfo VName ->+  TermM (ExpBase Info VName)+checkRecursive _ _ [] body = checkExp body+checkRecursive fname loc params' body = do+  ftype <- newType loc Lifted (baseName fname) NoUniqueness+  let bindF scope =+        scope {scopeVtable = M.insert fname (BoundV [] ftype) $ scopeVtable scope}+  body' <- localScope bindF $ checkExp body+  body_t <- expType body'+  let fun_t =+        foldFunType+          (map (first (const ()) . patternType) params')+          (RetType [] $ bimap (const ()) (const Nonunique) body_t)+  ctEq (Reason (locOf loc)) ftype fun_t+  pure body'++-- | Replace artificial variables with the types they denote, so that no+-- artificial variable leaks into the result.+onArtificial ::+  M.Map TyVar (TypeBase () NoUniqueness) ->+  M.Map TyVar (TypeBase () NoUniqueness) ->+  M.Map TyVar (TypeBase () NoUniqueness)+onArtificial artificial solution =+  M.map (substTyVars (`M.lookup` solution) . first (const ())) artificial+    <> solution++-- | Type check a single value definition.+checkValDef ::+  ( VName,+    Maybe (TypeExp (ExpBase NoInfo VName) VName),+    [TypeParam],+    [PatBase NoInfo VName ParamType],+    ExpBase NoInfo VName,+    SrcLoc+  ) ->+  TypeM+    ( Either TypeError ([TypeParam], M.Map TyVar (TypeBase () NoUniqueness)),+      [Pat ParamType],+      Maybe (TypeExp Exp VName),+      Exp+    )+checkValDef (fname, retdecl, tparams, params, body, loc) = runTermM $ do+  (params', body', retdecl') <-+    bindParams tparams params $ \params' -> do+      body' <- checkRecursive fname loc params' body+      (_, retdecl') <- checkRetDecl body' retdecl+      pure (params', body', retdecl')++  cts <- gets termConstraints+  tyvars <- gets termTyVars+  typarams <- gets termTyParams+  artificial <- gets $ M.map (first (const ())) . termArtificial++  solution <-+    bitraverse+      pure+      (fmap (second (onArtificial artificial)) . onTySolution params' body')+      $ solve (reverse cts) typarams tyvars+  pure (solution, params', retdecl', body')+  where+    onTySolution params' body' (unconstrained, solution) = do+      checkTyInstLiftedness solution+      body_t <- expType body'+      let fun_t =+            foldFunType+              (map (first (const ()) . patternType) params')+              (RetType [] $ bimap (const ()) (const Nonunique) body_t)+      generaliseAndDefaults unconstrained solution fun_t++-- | Type check a single expression, which may have a polymorphic+-- type.+checkSingleExp ::+  ExpBase NoInfo VName ->+  TypeM+    ( Either TypeError ([TypeParam], M.Map TyVar (TypeBase () NoUniqueness)),+      Exp+    )+checkSingleExp e = runTermM $ do+  e' <- checkExp e+  cts <- gets termConstraints+  tyvars <- gets termTyVars+  typarams <- gets termTyParams+  artificial <- gets termArtificial++  case solve cts typarams tyvars of+    Left err -> pure (Left err, e')+    Right (unconstrained, solution) -> do+      checkTyInstLiftedness solution+      e_t <- expType e'+      x <-+        second (onArtificial (M.map (first (const ())) artificial))+          <$> generaliseAndDefaults unconstrained solution (first (const ()) e_t)+      pure (Right x, e')++-- | Type-check a single size expression in isolation, which must have+-- type @i64@.+checkSizeExp ::+  ExpBase NoInfo VName ->+  TypeM+    ( Either TypeError ([UnconTyVar], M.Map TyVar (TypeBase () NoUniqueness)),+      Exp+    )+checkSizeExp e = runTermM $ do+  e' <- checkSizeExp' e+  cts <- gets termConstraints+  tyvars <- gets termTyVars+  typarams <- gets termTyParams+  artificial <- gets termArtificial++  case solve cts typarams tyvars of+    Left err -> pure (Left err, e')+    Right (unconstrained, solution) -> do+      checkTyInstLiftedness solution+      solution' <-+        onArtificial (M.map (first (const ())) artificial)+          <$> doDefaults mempty solution+      pure (Right (unconstrained, solution'), e')
+ src/Language/Futhark/TypeChecker/TySolve.hs view
@@ -0,0 +1,789 @@+module Language.Futhark.TypeChecker.TySolve+  ( Type,+    Solution,+    UnconTyVar,+    solve,+  )+where++import Control.Monad+import Control.Monad.Except+import Control.Monad.Reader+import Control.Monad.ST+import Data.Bifunctor+import Data.List qualified as L+import Data.Loc+import Data.Map qualified as M+import Data.Maybe+import Data.Set qualified as S+import Futhark.Util.Pretty+import Language.Futhark+import Language.Futhark.TypeChecker.Constraints+import Language.Futhark.TypeChecker.Error+import Language.Futhark.TypeChecker.Monad (Notes, TypeError (..), aNote, withIndexLink)+import Language.Futhark.TypeChecker.UnionFind++-- | The type representation used by the constraint solver. Agnostic+-- to sizes and uniqueness.+type Type = CtType ()++type UF s = M.Map TyVar (TyVarNode s)++newtype SolverState s = SolverState {solverTyVars :: UF s}++newtype SolveM s a = SolveM+  { runSolveM :: ExceptT TypeError (ReaderT (SolverState s) (ST s)) a+  }+  deriving (Functor, Applicative, Monad, MonadError TypeError, MonadReader (SolverState s))++-- | A solution maps a type variable to its substitution. This+-- substitution is complete, in the sense there are no right-hand+-- sides that contain a type variable.+type Solution = M.Map TyVar (Either [PrimType] (TypeBase () NoUniqueness))++-- | An unconstrained type variable comprises a name and (ironically)+-- a constraint on how it can be instantiated.+type UnconTyVar = (VName, Liftedness)++liftST :: ST s a -> SolveM s a+liftST = SolveM . lift . lift++getSol' :: TyVarNode s -> SolveM s TyVarSol+getSol' = liftST . getSol++union' :: TyVarNode s -> TyVarNode s -> SolveM s ()+union' tv1 tv2 = liftST $ union tv1 tv2++unionNewSol' tv1 tv2 new_sol = liftST $ unionNewSol tv1 tv2 new_sol++unionNewSol' :: TyVarNode s -> TyVarNode s -> TyVarSol -> SolveM s ()+getKey' :: TyVarNode s -> SolveM s TyVar+getKey' = liftST . getKey++initializeState :: TyParams -> TyVars () -> ST s (SolverState s)+initializeState typarams tyvars = do+  tyvars' <- M.traverseWithKey f tyvars+  typarams' <- M.traverseWithKey g typarams+  pure $ SolverState $ typarams' <> tyvars'+  where+    f tv (_lvl, info) = makeTyVarNode tv info+    g tv (lvl, lft, loc) = makeTyParamNode tv lvl lft loc++typeError :: Loc -> Notes -> Doc () -> SolveM s ()+typeError loc notes msg =+  throwError $ TypeError loc notes msg++typeVar :: (Monoid u) => VName -> TypeBase dim u+typeVar v = Scalar $ TypeVar mempty (qualName v) []++cannotUnify ::+  Reason Type ->+  Notes ->+  BreadCrumbs ->+  Type ->+  Type ->+  SolveM s ()+cannotUnify reason notes bcs t1 t2 = do+  t1' <- substTyVars t1+  t2' <- substTyVars t2+  case reason of+    Reason loc ->+      typeError loc notes . stack $+        [ "Cannot unify",+          indent 2 (pretty t1'),+          "with",+          indent 2 (pretty t2')+        ]+          <> [pretty bcs | not $ hasNoBreadCrumbs bcs]+    ReasonPatMatch loc pat value_t ->+      typeError loc notes . stack $+        [ "Pattern",+          indent 2 $ align $ pretty pat,+          "cannot match value of type",+          indent 2 $ align $ pretty value_t+        ]+          <> [pretty bcs | not $ hasNoBreadCrumbs bcs]+    ReasonAscription loc expected actual ->+      typeError loc notes . stack $+        [ "Expression does not have expected type from type ascription.",+          "Expected:" <+> align (pretty expected),+          "Actual:  " <+> align (pretty actual)+        ]+          <> [pretty bcs | not $ hasNoBreadCrumbs bcs]+    ReasonRetType loc expected actual -> do+      expected' <- substTyVars expected+      actual' <- substTyVars actual+      typeError loc notes . stack $+        [ "Function body does not have expected type.",+          "Expected:" <+> align (pretty expected'),+          "Actual:  " <+> align (pretty actual')+        ]+          <> [pretty bcs | not $ hasNoBreadCrumbs bcs]+    ReasonApply loc f e expected actual -> do+      expected' <- substTyVars expected+      actual' <- substTyVars actual+      typeError loc notes . stack $+        [ header,+          "Expected:" <+> align (pretty expected'),+          "Actual:  " <+> align (pretty actual')+        ]+      where+        header =+          case f of+            (Nothing, _) ->+              "Cannot apply function to"+                <+> dquotes (shorten $ group $ pretty e)+                <> " (invalid type)."+            (Just fname, _) ->+              "Cannot apply"+                <+> dquotes (pretty fname)+                <+> "to"+                <+> dquotes (align $ shorten $ group $ pretty e)+                <> " (invalid type)."+    ReasonApplySplit loc (fname, 0) _ ftype ->+      typeError loc notes $+        stack+          [ "Cannot apply"+              <+> fname'+              <+> "as function, as it has non-function type:"+              </> indent 2 (align $ pretty ftype)+          ]+      where+        fname' = maybe "expression" (dquotes . pretty) fname+    ReasonApplySplit loc (fname, i) e _ ->+      typeError loc notes $+        stack+          [ "Cannot apply"+              <+> fname'+              <+> "to"+              <+> dquotes (align $ shorten $ group $ pretty e)+              <> ".",+            "Function accepts only" <+> pretty i <+> "arguments."+          ]+      where+        fname' = maybe "expression" (dquotes . pretty) fname+    ReasonBranches loc former latter -> do+      former' <- substTyVars former+      latter' <- substTyVars latter+      typeError loc notes . stack $+        [ "Branches differ in type.",+          "Former:" <+> pretty former',+          "Latter:" <+> pretty latter'+        ]++unsharedConstructorsMsg :: M.Map Name t -> M.Map Name t -> Doc a+unsharedConstructorsMsg cs1 cs2 =+  "Unshared constructors:" <+> commasep (map (("#" <>) . pretty) missing) <> "."+  where+    missing =+      filter (`notElem` M.keys cs1) (M.keys cs2)+        ++ filter (`notElem` M.keys cs2) (M.keys cs1)++substTyVars :: (Monoid u) => TypeBase () u -> SolveM s (TypeBase () u)+substTyVars (Scalar (TypeVar u qn args)) = do+  mb_node <- maybeLookupUF $ qualLeaf qn+  case mb_node of+    Just node -> do+      sol <- getSol' node+      qn_k <- qualName <$> getKey' node+      case sol of+        Solved t -> do+          t' <- substTyVars t+          pure $ second (const mempty) t'+        _ -> makeTyVar qn_k+    _ -> makeTyVar qn+  where+    makeTyVar qn' = do+      args' <- mapM onArg args+      pure $ Scalar $ TypeVar u qn' args'+    onArg (TypeArgType t) = TypeArgType <$> substTyVars t+    onArg d@(TypeArgDim _) = pure d+substTyVars p@(Scalar (Prim _)) = pure p+substTyVars (Scalar (Record fs)) =+  Scalar . Record <$> traverse substTyVars fs+substTyVars (Scalar (Sum cs)) =+  Scalar . Sum <$> traverse (mapM substTyVars) cs+substTyVars (Scalar (Arrow u pname d t1 (RetType ext t2))) = do+  t1' <- substTyVars t1+  t2' <- substTyVars t2+  pure $+    Scalar $+      Arrow u pname d t1' $+        RetType ext $+          t2' `setUniqueness` uniqueness t2+substTyVars (Array u shape elemt) = do+  elemt' <- substTyVars $ Scalar elemt+  pure $ arrayOfWithAliases u shape elemt'++occursCheck :: Reason Type -> VName -> VName -> Type -> SolveM s ()+occursCheck reason v k tp = do+  let vars = typeVars tp+  when (k `S.member` vars)+    . typeError (locOf reason) mempty+    . withIndexLink "occurs-check"+    $ "Occurs check: cannot instantiate"+      <+> prettyName v+      <+> "with"+      <+> pretty tp+      <> "."++bindTyVar ::+  Reason Type ->+  BreadCrumbs ->+  VName ->+  TyVarNode s ->+  Type ->+  SolveM s ()+bindTyVar reason bcs v v_node t' = do+  t <- substTyVars t'+  k <- getKey' v_node+  occursCheck reason v k t++  v_info <- getSol' v_node++  setInfo v_node $ Solved t++  case (v_info, t) of+    (Unsolved TyVarFree {}, _) -> pure ()+    (Unsolved (TyVarPrim _ v_pts), _) ->+      if t `elem` map (Scalar . Prim) v_pts+        then pure ()+        else cannotUnify reason notes bcs (typeVar v) t+      where+        notes =+          aNote $+            "Cannot instantiate type that must be one of"+              </> indent 2 (pretty v_pts)+              </> "with"+              </> indent 2 (pretty t)+    (Unsolved (TyVarSum _ cs1), Scalar (Sum cs2)) ->+      if all (`elem` M.keys cs2) (M.keys cs1)+        then unifySharedConstructors reason bcs cs1 cs2+        else cannotUnify reason notes bcs (typeVar v) t+      where+        notes =+          aNote $+            "Cannot match type with constructors"+              </> indent 2 (stack (map (("#" <>) . pretty) (M.keys cs1)))+              </> "with type with constructors"+              </> indent 2 (stack (map (("#" <>) . pretty) (M.keys cs2)))+              </> unsharedConstructorsMsg cs1 cs2+    (Unsolved (TyVarSum _ cs1), _) ->+      typeError (locOf reason) mempty $+        "Cannot unify type with constructors"+          </> indent 2 (pretty (Sum cs1))+          </> "with type"+          </> indent 2 (pretty t)+    (Unsolved (TyVarRecord _ fs1), Scalar (Record fs2)) ->+      if all (`elem` M.keys fs2) (M.keys fs1)+        then unifySharedFields reason bcs fs1 fs2+        else+          typeError (locOf reason) mempty $+            "Cannot unify record type with fields"+              </> indent 2 (pretty (Record fs1))+              </> "with record type"+              </> indent 2 (pretty (Record fs2))+    (Unsolved (TyVarRecord _ fs1), _) ->+      typeError (locOf reason) mempty $+        "Cannot unify record type with fields"+          </> indent 2 (pretty (Record fs1))+          </> "with type"+          </> indent 2 (pretty t)+    --+    -- Internal error cases+    (Solved {}, _) ->+      error $ "Type variable already solved: " <> prettyNameString v+    (Param {}, _) ->+      error $ "Cannot substitute type parameter: " <> prettyNameString v++solveCt :: CtTy () -> SolveM s ()+solveCt (CtEq reason t1 t2) = solveEq reason mempty t1 t2++solveEq :: Reason Type -> BreadCrumbs -> Type -> Type -> SolveM s ()+solveEq reason obcs orig_t1 orig_t2 = do+  solveCt' (obcs, (orig_t1, orig_t2))+  where+    flexible :: VName -> SolveM s (Maybe (TyVarNode s))+    flexible v = do+      uf <- asks solverTyVars+      case M.lookup v uf of+        j_n@(Just node) -> do+          sol <- getSol' node+          pure $ case sol of+            Unsolved _ -> j_n+            _ -> Nothing+        Nothing -> pure Nothing++    normalize :: TypeBase () NoUniqueness -> SolveM s (TypeBase () NoUniqueness)+    normalize t@(Scalar (TypeVar _ (QualName [] v) [])) = do+      uf <- asks solverTyVars+      case M.lookup v uf of+        Just node -> do+          sol <- getSol' node+          case sol of+            Solved t' -> normalize t'+            _ -> typeVar <$> getKey' node+        Nothing -> pure t+    normalize t = pure t++    solveCt' :: (BreadCrumbs, (Type, Type)) -> SolveM s ()+    solveCt' (bcs, (t1, t2)) = do+      t1' <- normalize t1+      t2' <- normalize t2+      case (t1', t2') of+        ( Scalar (TypeVar _ (QualName [] v1) []),+          Scalar (TypeVar _ (QualName [] v2) [])+          )+            | v1 == v2 -> pure ()+            | otherwise -> do+                mb_node1 <- flexible v1+                mb_node2 <- flexible v2+                case (mb_node1, mb_node2) of+                  (Nothing, Nothing) ->+                    cannotUnify reason mempty bcs t1 t2+                  (Just v1_node, Nothing) ->+                    bindTyVar reason bcs v1 v1_node t2'+                  (Nothing, Just v2_node) ->+                    bindTyVar reason bcs v2 v2_node t1'+                  (Just v1_node, Just v2_node) ->+                    unionTyVars reason bcs v1 v1_node v2_node+        (Scalar (TypeVar _ (QualName [] v1) []), _) -> do+          mb_node <- flexible v1+          case mb_node of+            Just node -> bindTyVar reason bcs v1 node t2'+            Nothing -> tryUnify t1' t2' reason bcs+        (_, Scalar (TypeVar _ (QualName [] v2) [])) -> do+          mb_node <- flexible v2+          case mb_node of+            Just node -> bindTyVar reason bcs v2 node t1'+            Nothing -> tryUnify t1' t2' reason bcs+        (_, _) -> tryUnify t1' t2' reason bcs++    tryUnify :: Type -> Type -> Reason Type -> BreadCrumbs -> SolveM s ()+    tryUnify t1 t2 r bcs =+      case unify t1 t2 of+        Left details -> cannotUnify r (foldMap aNote details) bcs t1 t2+        Right eqs -> mapM_ solveCt' eqs++-- | Unify at the root, emitting new equalities that must hold.+unify :: Type -> Type -> Either (Maybe (Doc a)) [(BreadCrumbs, (Type, Type))]+unify (Scalar (Prim pt1)) (Scalar (Prim pt2))+  | pt1 == pt2 = Right []+unify+  (Scalar (TypeVar _ (QualName _ v1) targs1))+  (Scalar (TypeVar _ (QualName _ v2) targs2))+    | v1 == v2 =+        Right $ mapMaybe f $ zip targs1 targs2+    where+      f (TypeArgType t1, TypeArgType t2) = Just (mempty, (t1, t2))+      f _ = Nothing+unify+  (Scalar (Arrow _ _ _ t1a (RetType _ t1r)))+  (Scalar (Arrow _ _ _ t2a (RetType _ t2r))) =+    Right [(mempty, (t1a, t2a)), (mempty, (t1r', t2r'))]+    where+      t1r' = t1r `setUniqueness` NoUniqueness+      t2r' = t2r `setUniqueness` NoUniqueness+unify (Scalar (Record fs1)) (Scalar (Record fs2))+  | M.keys fs1 == M.keys fs2 =+      Right $+        map (first matchingField) $+          M.toList $+            M.intersectionWith (,) fs1 fs2+  | Just n1 <- length <$> areTupleFields fs1,+    Just n2 <- length <$> areTupleFields fs2,+    n1 /= n2 =+      Left . Just $+        "Tuples have"+          <+> pretty n1+          <+> "and"+          <+> pretty n2+          <+> "elements respectively."+  | otherwise =+      let missing =+            filter (`notElem` M.keys fs1) (M.keys fs2)+              <> filter (`notElem` M.keys fs2) (M.keys fs1)+       in Left . Just $+            "unshared fields:" <+> commasep (map pretty missing) <> "."+unify (Scalar (Sum cs1)) (Scalar (Sum cs2))+  | M.keys cs1 == M.keys cs2 =+      fmap concat . forM cs' $ \(c, (ts1, ts2)) -> do+        if length ts1 == length ts2+          then Right $ zipWith (curry (matchingConstructor c,)) ts1 ts2+          else Left Nothing+  | otherwise =+      Left . Just $ unsharedConstructorsMsg cs1 cs2+  where+    cs' = M.toList $ M.intersectionWith (,) cs1 cs2+unify t1 t2+  | Just t1' <- peelArray 1 t1,+    Just t2' <- peelArray 1 t2 =+      Right [(mempty, (t1', t2'))]+unify _ _ = Left Nothing++maybeLookupTyVarSol :: TyVar -> SolveM s (Maybe TyVarSol)+maybeLookupTyVarSol tv = do+  tyvars <- asks solverTyVars+  case M.lookup tv tyvars of+    Nothing -> pure Nothing+    Just node -> do+      sol <- getSol' node+      pure $ Just sol++lookupTyVar :: TyVar -> SolveM s (Either (TyVarInfo ()) Type)+lookupTyVar tv =+  maybe bad unpack <$> maybeLookupTyVarSol tv+  where+    bad = error $ "Unknown tyvar: " <> prettyNameString tv+    unpack (Param {}) = error $ "Is a type param: " <> prettyNameString tv+    unpack (Solved t) = Right t+    unpack (Unsolved info) = Left info++lookupTyVarInfo :: TyVarNode s -> SolveM s (TyVarInfo ())+lookupTyVarInfo v_node = do+  r <- getSol' v_node+  case r of+    Unsolved info -> pure info+    _ -> do+      v <- getKey' v_node+      error $ "Tyvar is nonflexible: " <> prettyNameString v++lookupUF :: TyVar -> SolveM s (TyVarNode s)+lookupUF tv = do+  uf <- asks solverTyVars+  case M.lookup tv uf of+    Nothing -> error $ "Unknown tyvar: " <> prettyNameString tv+    Just node -> pure node++unifySharedFields ::+  Reason Type ->+  BreadCrumbs ->+  M.Map Name Type ->+  M.Map Name Type ->+  SolveM s ()+unifySharedFields reason bcs fs1 fs2 =+  forM_ (M.toList $ M.intersectionWith (,) fs1 fs2) $ \(f, (ts1, ts2)) ->+    solveEq reason (matchingField f <> bcs) ts1 ts2++unifySharedConstructors ::+  Reason Type ->+  BreadCrumbs ->+  M.Map Name [Type] ->+  M.Map Name [Type] ->+  SolveM s ()+unifySharedConstructors reason bcs cs1 cs2 =+  forM_ (M.toList $ M.intersectionWith (,) cs1 cs2) $ \(c, (ts1, ts2)) ->+    if length ts1 == length ts2+      then zipWithM_ (solveEq reason $ matchingConstructor c <> bcs) ts1 ts2+      else+        typeError (locOf reason) mempty $+          "Cannot unify type with constructor"+            </> indent 2 (pretty (Sum (M.singleton c ts1)))+            </> "with type of constructor"+            </> indent 2 (pretty (Sum (M.singleton c ts2)))+            </> "because they differ in arity."++setInfo :: TyVarNode s -> TyVarSol -> SolveM s ()+setInfo node sol = liftST $ assignNewSol node sol++unionTyVars ::+  Reason Type ->+  BreadCrumbs ->+  VName ->+  TyVarNode s ->+  TyVarNode s ->+  SolveM s ()+unionTyVars reason bcs v v_node t_node = do+  v_sol <- getSol' v_node+  t_info <- lookupTyVarInfo t_node+  c <- check v_sol t_info+  case c of+    Left (loc, notes, msg) -> typeError loc notes msg+    Right (Just new_sol) -> unionNewSol' v_node t_node new_sol+    Right Nothing -> union' v_node t_node+  where+    check ::+      TyVarSol ->+      TyVarInfo () ->+      SolveM s (Either (Loc, Notes, Doc ()) (Maybe TyVarSol))+    check v_sol t_info =+      case (v_sol, t_info) of+        (Unsolved (TyVarFree _ v_l), TyVarFree t_loc t_l)+          | v_l /= t_l ->+              pure $ Right $ Just $ Unsolved $ TyVarFree t_loc (min v_l t_l)+        (Unsolved info, TyVarFree {}) -> do+          pure $ Right $ Just $ Unsolved info+        --+        -- TyVarPrim cases+        ( Unsolved (TyVarPrim _ v_pts),+          TyVarPrim t_loc t_pts+          ) ->+            let pts = L.intersect v_pts t_pts+             in case pts of+                  [] ->+                    pure $+                      Left+                        ( locOf reason,+                          mempty,+                          "Cannot unify type that must be one of"+                            </> indent 2 (pretty v_pts)+                            </> "with type that must be one of"+                            </> indent 2 (pretty t_pts)+                        )+                  _ -> pure $ Right $ Just $ Unsolved $ TyVarPrim t_loc pts+        (Unsolved (TyVarPrim _ v_pts), TyVarRecord {}) ->+          pure $+            Left+              ( locOf reason,+                mempty,+                "Cannot unify type that must be one of"+                  </> indent 2 (pretty v_pts)+                  </> "with type that must be a record."+              )+        (Unsolved (TyVarPrim _ v_pts), TyVarSum {}) ->+          pure $+            Left+              ( locOf reason,+                mempty,+                "Cannot unify type that must be one of"+                  </> indent 2 (pretty v_pts)+                  </> "with type that must be sum."+              )+        --+        -- TyVarSum cases+        ( Unsolved (TyVarSum _ cs1),+          TyVarSum loc cs2+          ) -> do+            unifySharedConstructors reason bcs cs1 cs2+            let cs3 = cs1 <> cs2+            pure $ Right $ Just $ Unsolved $ TyVarSum loc cs3+        ( Unsolved TyVarSum {},+          TyVarPrim _ pts+          ) ->+            pure $+              Left+                ( locOf reason,+                  mempty,+                  "A sum type cannot be one of"+                    </> indent 2 (pretty pts)+                )+        ( Unsolved (TyVarSum _ cs1),+          TyVarRecord _ fs+          ) ->+            pure $+              Left+                ( locOf reason,+                  mempty,+                  "Cannot unify type with constructors"+                    </> indent 2 (pretty (Sum cs1))+                    </> "with type"+                    </> indent 2 (pretty (Scalar (Record fs)))+                )+        --+        -- TyVarRecord cases+        ( Unsolved (TyVarRecord _ fs1),+          TyVarRecord loc fs2+          ) -> do+            unifySharedFields reason bcs fs1 fs2+            let fs3 = fs1 <> fs2+            pure $ Right $ Just $ Unsolved $ TyVarRecord loc fs3+        ( Unsolved TyVarRecord {},+          TyVarPrim _ pts+          ) ->+            pure $+              Left+                ( locOf reason,+                  mempty,+                  "A record type cannot be one of"+                    </> indent 2 (pretty pts)+                )+        ( Unsolved (TyVarRecord _ fs1),+          TyVarSum _ cs+          ) ->+            pure $+              Left+                ( locOf reason,+                  mempty,+                  "Cannot unify record type"+                    </> indent 2 (pretty (Record fs1))+                    </> "with type"+                    </> indent 2 (pretty (Scalar (Sum cs)))+                )+        --+        -- Internal error cases+        (Solved {}, _) -> alreadySolved+        (Param {}, _) -> isParam+        _ -> pure $ Right Nothing++    alreadySolved = error $ "Type variable already solved: " <> prettyNameString v+    isParam = error $ "Type name is a type parameter: " <> prettyNameString v++scopeViolation :: Reason Type -> VName -> Type -> VName -> SolveM s ()+scopeViolation reason v1 ty v2 =+  typeError (locOf reason) mempty . withIndexLink "scope-violation" $+    "Cannot unify type"+      </> indent 2 (pretty ty)+      </> "with"+      <+> dquotes (prettyName v1)+      <+> "(scope violation)."+      </> "This is because"+      <+> dquotes (prettyName v2)+      <+> "is rigidly bound in a deeper scope."++scopeCheck :: Reason Type -> TyVar -> Level -> Type -> SolveM s ()+scopeCheck reason v v_lvl ty = mapM_ check $ typeVars ty+  where+    check :: TyVar -> SolveM s ()+    check ty_v = do+      maybe (pure ()) checkNode =<< maybeLookupUF ty_v++    checkNode :: TyVarNode s -> SolveM s ()+    checkNode node = do+      sol <- getSol' node+      case sol of+        Param ty_v_lvl _ _+          | ty_v_lvl > v_lvl -> do+              k <- getKey' node+              ty' <- substTyVars ty+              scopeViolation reason v ty' k+        Solved ty' -> do+          mapM_ check $ typeVars ty'+        _ -> pure ()++-- | If a type variable has a liftedness constraint, we propagate that+-- constraint to its solution. The actual checking for correct usage+-- is done later, by 'localChecks' and 'instTyVars' in the sized type+-- checker, which know why the constraints exist and can produce+-- proper error messages.+liftednessCheck :: Liftedness -> Type -> SolveM s ()+liftednessCheck l (Scalar (TypeVar _ (QualName [] v) _)) = do+  v_info <- maybeLookupTyVarSol v+  case v_info of+    Nothing ->+      -- Is an opaque type.+      pure ()+    Just (Solved v_ty) ->+      liftednessCheck l v_ty+    Just Param {} -> pure ()+    Just (Unsolved (TyVarFree loc v_l))+      | l < v_l -> do+          node <- lookupUF v+          setInfo node $ Unsolved $ TyVarFree loc l+    Just Unsolved {} -> pure ()+liftednessCheck _ (Scalar Prim {}) = pure ()+liftednessCheck Lifted _ = pure ()+liftednessCheck _ Array {} = pure ()+liftednessCheck _ (Scalar Arrow {}) = pure ()+liftednessCheck l (Scalar (Record fs)) =+  mapM_ (liftednessCheck l) fs+liftednessCheck l (Scalar (Sum cs)) =+  mapM_ (mapM_ $ liftednessCheck l) cs+liftednessCheck _ (Scalar TypeVar {}) = pure ()++solveTyVar :: (VName, (Level, TyVarInfo ())) -> SolveM s ()+solveTyVar (tv, (lvl, TyVarFree loc l)) = do+  tv_t <- lookupTyVar tv+  case tv_t of+    Right ty -> do+      scopeCheck (Reason loc) tv lvl ty+      liftednessCheck l ty+    _ -> pure ()+solveTyVar (tv, (_, TyVarPrim loc pts)) = do+  tv_t <- lookupTyVar tv+  case tv_t of+    Right ty+      | ty `elem` map (Scalar . Prim) pts -> pure ()+      | otherwise ->+          typeError loc mempty $+            "Numeric constant inferred to be of type"+              </> indent 2 (align (pretty ty))+              </> "which is not possible."+    _ -> pure ()+solveTyVar (tv, (_, TyVarRecord loc fs1)) = do+  tv_t <- lookupTyVar tv+  case tv_t of+    Left _ ->+      typeError loc mempty . withIndexLink "ambiguous-type" $+        "Type"+          <+> prettyName tv+          <+> "is ambiguous."+          </> "Must be a record with fields"+          </> indent 2 (pretty (Scalar (Record fs1)))+    Right _ -> pure ()+solveTyVar (tv, (_, TyVarSum loc cs1)) = do+  tv_t <- lookupTyVar tv+  case tv_t of+    Left _ ->+      typeError loc mempty . withIndexLink "ambiguous-type" $+        "Type is ambiguous."+          </> "Must be a sum type with constructors"+          </> indent 2 (pretty (Scalar (Sum cs1)))+    Right _ -> pure ()++maybeLookupUF :: TyVar -> SolveM s (Maybe (TyVarNode s))+maybeLookupUF tv = do+  uf <- asks solverTyVars+  pure . M.lookup tv $ uf++getSolution :: SolveM s ([UnconTyVar], Solution)+getSolution = do+  uf <- asks solverTyVars+  resolved <- M.traverseWithKey resolve uf+  let unconstrained = M.foldrWithKey unconstr [] resolved+      sol = M.mapMaybeWithKey mkSubst resolved+  pure (unconstrained, sol)+  where+    resolve ::+      TyVar ->+      TyVarNode s ->+      SolveM s (Either [PrimType] (TypeBase () NoUniqueness), Maybe Liftedness)+    resolve tv node = do+      sol <- getSol' node+      case sol of+        Unsolved (TyVarFree _ l) -> do+          k <- getKey' node+          let tv' = typeVar k+          -- If the current type variable and root type variable are+          -- different, this variable is unconstrained, so we save the+          -- liftedness constraint for later.+          pure (Right tv', if k == tv then Just l else Nothing)+        Unsolved (TyVarPrim _ pts) -> pure (Left pts, Nothing)+        Solved t -> do+          t' <- substTyVars t+          pure (Right $ first (const ()) t', Nothing)+        _ -> do+          k <- getKey' node+          pure (Right $ typeVar k, Nothing)++    unconstr ::+      TyVar ->+      (Either [PrimType] (TypeBase () NoUniqueness), Maybe Liftedness) ->+      [UnconTyVar] ->+      [UnconTyVar]+    unconstr tv (_, Just l) acc = (tv, l) : acc+    unconstr _ _ acc = acc++    mkSubst ::+      TyVar ->+      (Either [PrimType] (TypeBase () NoUniqueness), Maybe Liftedness) ->+      Maybe (Either [PrimType] (TypeBase () NoUniqueness))+    mkSubst _ (_, Just _) = Nothing+    mkSubst tv (s@(Right (Scalar (TypeVar _ (QualName [] tv') _))), _) =+      if tv /= tv' then Just s else Nothing+    mkSubst _ (s, _) = Just s++-- | Solve type constraints, producing either an error or a solution,+-- alongside a list of unconstrained type variables.+solve ::+  [CtTy ()] ->+  TyParams ->+  TyVars () ->+  Either TypeError ([UnconTyVar], Solution)+solve constraints typarams tyvars =+  runST $ do+    r <- initializeState typarams tyvars+    flip runReaderT r $ runExceptT $ runSolveM $ do+      mapM_ solveCt constraints+      mapM_ solveTyVar $ M.toList tyvars+      getSolution+{-# NOINLINE solve #-}
src/Language/Futhark/TypeChecker/Types.hs view
@@ -8,6 +8,7 @@     TypeSubs,     Substitutable (..),     substTypesAny,+    substTyVars,      -- * Witnesses     mustBeExplicitInType,@@ -22,6 +23,7 @@ import Data.Bifunctor import Data.List qualified as L import Data.Map.Strict qualified as M+import Data.Maybe (fromMaybe, isNothing) import Data.Set qualified as S import Futhark.Util (nubOrd) import Futhark.Util.Pretty@@ -59,7 +61,7 @@ mustBeExplicitInBinding bind_t =   let (ts, ret) = unfoldFunType bind_t       alsoRet = M.unionWith (&&) $ M.fromList $ map (,True) (S.toList (fvVars (freeInType ret)))-   in S.fromList $ M.keys $ M.filter id $ alsoRet $ L.foldl' onType mempty $ map toStruct ts+   in S.fromList $ M.keys $ M.filter id $ alsoRet $ L.foldl' onType mempty $ map (toStruct . snd) ts   where     onType uses t = uses <> mustBeExplicitAux t -- Left-biased union. @@ -442,80 +444,128 @@     mkSubst p a =       error $ "applyType mkSubst: cannot substitute " ++ prettyString a ++ " for " ++ prettyString p +-- In case we are substituting the same RetType in multiple+-- places, we must ensure each instance is given distinct+-- dimensions.  E.g. substituting 'a ↦ ?[n].[n]bool' into '(a,a)'+-- should give '?[n][m].([n]bool,[m]bool)'.+--+-- XXX: the size names we invent here not globally unique.  This+-- is _probably_ not a problem, since substituting types with+-- outermost non-null existential sizes is done only when type+-- checking modules and monomorphising.+freshDims ::+  (Monoid as) =>+  RetTypeBase Size as ->+  State [VName] (RetTypeBase Size as)+freshDims (RetType [] t) = pure $ RetType [] t+freshDims (RetType ext t) = do+  seen_ext <- get+  if not $ any (`elem` seen_ext) ext+    then pure $ RetType ext t+    else do+      let start = maximum $ map baseTag seen_ext+          ext' = zipWith VName (map baseName ext) [start + 1 ..]+          mkSubst = ExpSubst . flip sizeFromName mempty . qualName+          extsubsts = M.fromList $ zip ext $ map mkSubst ext'+          RetType [] t' = substTypesRet (`M.lookup` extsubsts) t+      pure $ RetType ext' t'+ substTypesRet ::   (Monoid u) =>   (VName -> Maybe (Subst (RetTypeBase Size u))) ->   TypeBase Size u ->   RetTypeBase Size u substTypesRet lookupSubst ot =-  uncurry (flip RetType) $ runState (onType ot) []+  let (t', dims) = runState (onType ot) []+   in RetType dims (fromMaybe ot t')   where-    -- In case we are substituting the same RetType in multiple-    -- places, we must ensure each instance is given distinct-    -- dimensions.  E.g. substituting 'a ↦ ?[n].[n]bool' into '(a,a)'-    -- should give '?[n][m].([n]bool,[m]bool)'.-    ---    -- XXX: the size names we invent here not globally unique.  This-    -- is _probably_ not a problem, since substituting types with-    -- outermost non-null existential sizes is done only when type-    -- checking modules and monomorphising.-    freshDims (RetType [] t) = pure $ RetType [] t-    freshDims (RetType ext t) = do-      seen_ext <- get-      if not $ any (`elem` seen_ext) ext-        then pure $ RetType ext t-        else do-          let start = maximum $ map baseTag seen_ext-              ext' = zipWith VName (map baseName ext) [start + 1 ..]-              mkSubst = ExpSubst . flip sizeFromName mempty . qualName-              extsubsts = M.fromList $ zip ext $ map mkSubst ext'-              RetType [] t' = substTypesRet (`M.lookup` extsubsts) t-          pure $ RetType ext' t'-+    -- 'onType' returns 'Nothing' when the substitution does not change+    -- the type, so that the (very common) unchanged parts are shared+    -- rather than reconstructed. 'fromMaybe' at each level splices in+    -- the original subterm for unchanged children.     onType ::       forall as.       (Monoid as) =>       TypeBase Size as ->-      State [VName] (TypeBase Size as)+      State [VName] (Maybe (TypeBase Size as)) -    onType (Array u shape et) =-      arrayOfWithAliases u (applySubst lookupSubst' shape)-        <$> onType (Scalar et)-    onType (Scalar (Prim t)) = pure $ Scalar $ Prim t+    onType (Array u shape et) = do+      et' <- onType (Scalar et)+      let shape' = onShape shape+      pure $ case (shape', et') of+        (Nothing, Nothing) -> Nothing+        _ ->+          Just $+            arrayOfWithAliases u (fromMaybe shape shape') (fromMaybe (Scalar et) et')+    onType (Scalar (Prim _)) = pure Nothing     onType (Scalar (TypeVar u v targs)) = do       targs' <- mapM subsTypeArg targs       case lookupSubst $ qualLeaf v of         Just (Subst ps rt) -> do           RetType ext t <- freshDims rt           modify (ext ++)-          pure $ second (<> u) $ applyType ps (second (const u) t) targs'+          let targs'' = zipWith fromMaybe targs targs'+          pure $ Just $ second (<> u) $ applyType ps (second (const u) t) targs''         _ ->-          pure $ Scalar $ TypeVar u v targs'-    onType (Scalar (Record ts)) =-      Scalar . Record <$> traverse onType ts-    onType (Scalar (Arrow u v d t1 t2)) =-      Scalar <$> (Arrow u v d <$> onType t1 <*> onRetType t2)-    onType (Scalar (Sum ts)) =-      Scalar . Sum <$> traverse (traverse onType) ts+          pure $+            if all isNothing targs'+              then Nothing+              else Just $ Scalar $ TypeVar u v (zipWith fromMaybe targs targs')+    onType (Scalar (Record ts)) = do+      ts' <- traverse onType ts+      pure $+        if all isNothing ts'+          then Nothing+          else Just $ Scalar $ Record $ M.intersectionWith fromMaybe ts ts'+    onType (Scalar (Arrow u v d t1 t2)) = do+      t1' <- onType t1+      t2' <- onRetType t2+      pure $ case (t1', t2') of+        (Nothing, Nothing) -> Nothing+        _ -> Just $ Scalar $ Arrow u v d (fromMaybe t1 t1') (fromMaybe t2 t2')+    onType (Scalar (Sum ts)) = do+      ts' <- traverse (traverse onType) ts+      pure $+        if all (all isNothing) ts'+          then Nothing+          else Just $ Scalar $ Sum $ M.intersectionWith (zipWith fromMaybe) ts ts'      onRetType (RetType dims t) = do       ext <- get-      let (t', ext') = runState (onType t) ext-          new_ext = ext' L.\\ ext-      case t of-        Scalar Arrow {} -> do-          put ext'-          pure $ RetType dims t'-        _ ->-          pure $ RetType (new_ext <> dims) t'+      case runState (onType t) ext of+        (Nothing, _) -> pure Nothing+        (Just t', ext') -> do+          let new_ext = ext' L.\\ ext+          case t of+            Scalar Arrow {} -> do+              put ext'+              pure $ Just $ RetType dims t'+            _ ->+              pure $ Just $ RetType (new_ext <> dims) t' +    -- Applied type arguments are rare; substitute conservatively (do+    -- not bother sharing them).     subsTypeArg (TypeArgType t) = do       let RetType dims t' = substTypesRet lookupSubst' t       modify (dims ++)-      pure $ TypeArgType t'+      pure $ Just $ TypeArgType t'     subsTypeArg (TypeArgDim v) =-      pure $ TypeArgDim $ applySubst lookupSubst' v+      pure $ TypeArgDim <$> onSize v +    onShape (Shape ds) =+      let ds' = map onSize ds+       in if all isNothing ds'+            then Nothing+            else Just $ Shape $ zipWith fromMaybe ds ds'++    -- A bare size variable is either substituted or unchanged (its own+    -- type is always i64). Anything more complex is substituted+    -- conservatively (rare in types).+    onSize (Var (QualName _ v) _ _)+      | Just (ExpSubst e') <- lookupSubst' v = Just e'+    onSize (Var {}) = Nothing+    onSize e = Just $ applySubst lookupSubst' e+     lookupSubst' = fmap (fmap $ second (const NoUniqueness)) . lookupSubst  -- | Perform substitutions, from type names to types, on a type. Works@@ -538,6 +588,26 @@                 anySize (baseTag (qualLeaf v))           toAny d = d        in first toAny ot'++-- | Substitution without caring about sizes.+substTyVars :: (Monoid u) => (VName -> Maybe (TypeBase d NoUniqueness)) -> TypeBase d u -> TypeBase d u+substTyVars f (Scalar (TypeVar u qn args)) =+  case f $ qualLeaf qn of+    Just t' -> second (const mempty) $ substTyVars f t'+    Nothing -> Scalar (TypeVar u qn (map onArg args))+      where+        onArg (TypeArgType t) = TypeArgType $ substTyVars f t+        onArg (TypeArgDim e) = TypeArgDim e+substTyVars _ (Scalar (Prim pt)) = Scalar $ Prim pt+substTyVars f (Scalar (Record fs)) = Scalar $ Record $ M.map (substTyVars f) fs+substTyVars f (Scalar (Sum cs)) = Scalar $ Sum $ M.map (map $ substTyVars f) cs+substTyVars f (Scalar (Arrow u pname d t1 (RetType ext t2))) =+  Scalar $+    Arrow u pname d (substTyVars f t1) $+      RetType ext $+        substTyVars f t2 `setUniqueness` uniqueness t2+substTyVars f (Array u shape elemt) =+  arrayOfWithAliases u shape $ substTyVars f $ Scalar elemt  -- Note [AnySize] --
src/Language/Futhark/TypeChecker/Unify.hs view
@@ -1,1353 +1,901 @@--- | Implementation of unification and other core type system building--- blocks.-module Language.Futhark.TypeChecker.Unify-  ( Constraint (..),-    Usage (..),-    mkUsage,-    mkUsage',-    Level,-    Constraints,-    MonadUnify (..),-    Rigidity (..),-    RigidSource (..),-    BreadCrumbs,-    sizeFree,-    noBreadCrumbs,-    hasNoBreadCrumbs,-    dimNotes,-    zeroOrderType,-    arrayElemType,-    mustHaveConstr,-    mustHaveField,-    mustBeOneOf,-    equalityType,-    normType,-    normTypeFully,-    unify,-    unifyMostCommon,-    doUnification,-  )-where--import Control.Monad-import Control.Monad.Except-import Control.Monad.Identity-import Control.Monad.Reader-import Control.Monad.State-import Data.List qualified as L-import Data.Map.Strict qualified as M-import Data.Maybe-import Data.Set qualified as S-import Data.Text qualified as T-import Futhark.Util (topologicalSort)-import Futhark.Util.Pretty-import Language.Futhark-import Language.Futhark.Traversals-import Language.Futhark.TypeChecker.Monad hiding (BoundV)-import Language.Futhark.TypeChecker.Types---- | A piece of information that describes what process the type--- checker currently performing.  This is used to give better error--- messages for unification errors.-data BreadCrumb-  = MatchingTypes StructType StructType-  | MatchingFields [Name]-  | MatchingConstructor Name-  | Matching (Doc ())--instance Pretty BreadCrumb where-  pretty (MatchingTypes t1 t2) =-    "When matching type"-      </> indent 2 (pretty t1)-      </> "with"-      </> indent 2 (pretty t2)-  pretty (MatchingFields fields) =-    "When matching types of record field"-      <+> dquotes (mconcat $ punctuate "." $ map pretty fields)-      <> dot-  pretty (MatchingConstructor c) =-    "When matching types of constructor" <+> dquotes (pretty c) <> dot-  pretty (Matching s) =-    unAnnotate s---- | Unification failures can occur deep down inside complicated types--- (consider nested records).  We leave breadcrumbs behind us so we--- can report the path we took to find the mismatch.-newtype BreadCrumbs = BreadCrumbs [BreadCrumb]---- | An empty path.-noBreadCrumbs :: BreadCrumbs-noBreadCrumbs = BreadCrumbs []---- | Is the path empty?-hasNoBreadCrumbs :: BreadCrumbs -> Bool-hasNoBreadCrumbs (BreadCrumbs xs) = null xs---- | Drop a breadcrumb on the path behind you.-breadCrumb :: BreadCrumb -> BreadCrumbs -> BreadCrumbs-breadCrumb (MatchingFields xs) (BreadCrumbs (MatchingFields ys : bcs)) =-  BreadCrumbs $ MatchingFields (ys ++ xs) : bcs-breadCrumb bc (BreadCrumbs bcs) =-  BreadCrumbs $ bc : bcs--instance Pretty BreadCrumbs where-  pretty (BreadCrumbs []) = mempty-  pretty (BreadCrumbs bcs) = line <> stack (map pretty bcs)---- | A usage that caused a type constraint.-data Usage = Usage (Maybe T.Text) Loc-  deriving (Show)---- | Construct a 'Usage' from a location and a description.-mkUsage :: (Located a) => a -> T.Text -> Usage-mkUsage = flip (Usage . Just) . locOf---- | Construct a 'Usage' that has just a location, but no particular--- description.-mkUsage' :: (Located a) => a -> Usage-mkUsage' = Usage Nothing . locOf--instance Pretty Usage where-  pretty (Usage Nothing loc) = "use at " <> textwrap (locText loc)-  pretty (Usage (Just s) loc) = textwrap s <+> "at" <+> textwrap (locText loc)--instance Located Usage where-  locOf (Usage _ loc) = locOf loc---- | The level at which a type variable is bound.  Higher means--- deeper.  We can only unify a type variable at level @i@ with a type--- @t@ if all type names that occur in @t@ are at most at level @i@.-type Level = Int---- | A constraint on a yet-ambiguous type variable.-data Constraint-  = NoConstraint Liftedness Usage-  | ParamType Liftedness Loc-  | Constraint StructRetType Usage-  | Overloaded [PrimType] Usage-  | HasFields Liftedness (M.Map Name StructType) Usage-  | Equality Usage-  | HasConstrs Liftedness (M.Map Name [StructType]) Usage-  | ParamSize Loc-  | -- | Is not actually a type, but a term-level size,-    -- possibly already set to something specific.-    Size (Maybe Exp) Usage-  | -- | A size that does not unify with anything --    -- created from the result of applying a function-    -- whose return size is existential, or otherwise-    -- hiding a size.-    UnknownSize Loc RigidSource-  deriving (Show)--instance Located Constraint where-  locOf (NoConstraint _ usage) = locOf usage-  locOf (ParamType _ usage) = locOf usage-  locOf (Constraint _ usage) = locOf usage-  locOf (Overloaded _ usage) = locOf usage-  locOf (HasFields _ _ usage) = locOf usage-  locOf (Equality usage) = locOf usage-  locOf (HasConstrs _ _ usage) = locOf usage-  locOf (ParamSize loc) = locOf loc-  locOf (Size _ usage) = locOf usage-  locOf (UnknownSize loc _) = locOf loc---- | Mapping from fresh type variables, instantiated from the type--- schemes of polymorphic functions, to (possibly) specific types as--- determined on application and the location of that application, or--- a partial constraint on their type.-type Constraints = M.Map VName (Level, Constraint)--lookupSubst :: VName -> Constraints -> Maybe (Subst StructRetType)-lookupSubst v constraints = case snd <$> M.lookup v constraints of-  Just (Constraint t _) -> Just $ Subst [] $ applySubst (`lookupSubst` constraints) t-  Just (Size (Just d) _) ->-    Just $ ExpSubst $ applySubst (`lookupSubst` constraints) d-  _ -> Nothing---- | The source of a rigid size.-data RigidSource-  = -- | A function argument that is not a constant or variable name.-    RigidArg (Maybe (QualName VName)) T.Text-  | -- | An existential return size.-    RigidRet (Maybe (QualName VName))-  | -- | Similarly to 'RigidRet', but produce by a loop.-    RigidLoop-  | -- | Produced by a complicated slice expression.-    RigidSlice (Maybe Size) T.Text-  | -- | Produced by a complicated range expression.-    RigidRange-  | -- | Mismatch in branches.-    RigidCond StructType StructType-  | -- | Invented during unification.-    RigidUnify-  | -- | A name used in a size went out of scope.-    RigidOutOfScope Loc VName-  deriving (Eq, Ord, Show)---- | The ridigity of a size variable.  All rigid sizes are tagged with--- information about how they were generated.-data Rigidity = Rigid RigidSource | Nonrigid-  deriving (Eq, Ord, Show)--prettySource :: Loc -> Loc -> RigidSource -> Doc ()-prettySource ctx loc (RigidRet Nothing) =-  "is unknown size returned by function at"-    <+> pretty (locStrRel ctx loc)-    <> "."-prettySource ctx loc (RigidRet (Just fname)) =-  "is unknown size returned by"-    <+> dquotes (pretty fname)-    <+> "at"-    <+> pretty (locStrRel ctx loc)-    <> "."-prettySource ctx loc (RigidArg fname arg) =-  "is value of argument"-    </> indent 2 (shorten (pretty arg))-    </> "passed to"-    <+> fname'-    <+> "at"-    <+> pretty (locStrRel ctx loc)-    <> "."-  where-    fname' = maybe "function" (dquotes . pretty) fname-prettySource ctx loc (RigidSlice d slice) =-  "is size produced by slice"-    </> indent 2 (shorten (pretty slice))-    </> d_desc-    <> "at"-      <+> pretty (locStrRel ctx loc)-    <> "."-  where-    d_desc = case d of-      Just d' -> "of dimension of size " <> dquotes (pretty d') <> " "-      Nothing -> mempty-prettySource ctx loc RigidLoop =-  "is unknown size of value returned at" <+> pretty (locStrRel ctx loc) <> "."-prettySource ctx loc RigidRange =-  "is unknown length of range at" <+> pretty (locStrRel ctx loc) <> "."-prettySource ctx loc (RigidOutOfScope boundloc v) =-  "is an unknown size arising from "-    <> dquotes (prettyName v)-    <> " going out of scope at "-    <> pretty (locStrRel ctx loc)-    <> "."-      </> "Originally bound at "-    <> pretty (locStrRel ctx boundloc)-    <> "."-prettySource _ _ RigidUnify =-  "is an artificial size invented during unification of functions with anonymous sizes."-prettySource ctx loc (RigidCond t1 t2) =-  "is unknown due to conditional expression at "-    <> pretty (locStrRel ctx loc)-    <> "."-      </> "One branch returns array of type: "-    <> align (pretty t1)-      </> "The other an array of type:       "-    <> align (pretty t2)---- | Retrieve notes describing the purpose or origin of the given--- t'Size'.  The location is used as the *current* location, for the--- purpose of reporting relative locations.-dimNotes :: (Located a, MonadUnify m) => a -> Exp -> m Notes-dimNotes ctx (Var d _ _) = do-  c <- M.lookup (qualLeaf d) <$> getConstraints-  case c of-    Just (_, UnknownSize loc rsrc) ->-      pure . aNote $-        dquotes (pretty d) <+> prettySource (locOf ctx) loc rsrc-    _ -> pure mempty-dimNotes _ _ = pure mempty--typeNotes :: (Located a, MonadUnify m) => a -> StructType -> m Notes-typeNotes ctx =-  fmap mconcat-    . mapM (dimNotes ctx . flip sizeFromName mempty . qualName)-    . S.toList-    . fvVars-    . freeInType--typeVarNotes :: (MonadUnify m) => VName -> m Notes-typeVarNotes v = maybe mempty (note . snd) . M.lookup v <$> getConstraints-  where-    note (HasConstrs _ cs _) =-      aNote $-        prettyName v-          <+> "="-          <+> hsep (map ppConstr (M.toList cs))-          <+> "..."-    note (Overloaded ts _) =-      aNote $ prettyName v <+> "must be one of" <+> mconcat (punctuate ", " (map pretty ts))-    note (HasFields _ fs _) =-      aNote $-        prettyName v-          <+> "="-          <+> braces (mconcat (punctuate ", " (map ppField (M.toList fs))))-    note _ = mempty--    ppConstr (c, _) = "#" <> pretty c <+> "..." <+> "|"-    ppField (f, _) = prettyName f <> ":" <+> "..."---- | Monads that which to perform unification must implement this type--- class.-class (Monad m) => MonadUnify m where-  getConstraints :: m Constraints-  putConstraints :: Constraints -> m ()-  modifyConstraints :: (Constraints -> Constraints) -> m ()-  modifyConstraints f = do-    x <- getConstraints-    putConstraints $ f x--  newTypeVar :: (Monoid als, Located a) => a -> Name -> m (TypeBase dim als)-  newDimVar :: Usage -> Rigidity -> Name -> m VName-  newRigidDim :: (Located a) => a -> RigidSource -> Name -> m VName-  newRigidDim loc = newDimVar (mkUsage' loc) . Rigid-  newFlexibleDim :: Usage -> Name -> m VName-  newFlexibleDim usage = newDimVar usage Nonrigid--  curLevel :: m Level--  matchError ::-    (Located loc) =>-    loc ->-    Notes ->-    BreadCrumbs ->-    StructType ->-    StructType ->-    m a--  unifyError ::-    (Located loc) =>-    loc ->-    Notes ->-    BreadCrumbs ->-    Doc () ->-    m a---- | Replace all type variables with their substitution.-normTypeFully :: (Substitutable a, MonadUnify m) => a -> m a-normTypeFully t = do-  constraints <- getConstraints-  pure $ applySubst (`lookupSubst` constraints) t---- | Replace any top-level type variable with its substitution.-normType :: (MonadUnify m) => StructType -> m StructType-normType t@(Scalar (TypeVar _ (QualName [] v) [])) = do-  constraints <- getConstraints-  case snd <$> M.lookup v constraints of-    Just (Constraint (RetType [] t') _) -> normType t'-    _ -> pure t-normType t = pure t--rigidConstraint :: Constraint -> Bool-rigidConstraint ParamType {} = True-rigidConstraint ParamSize {} = True-rigidConstraint UnknownSize {} = True-rigidConstraint _ = False--unsharedConstructorsMsg :: M.Map Name t -> M.Map Name t -> Doc a-unsharedConstructorsMsg cs1 cs2 =-  "Unshared constructors:" <+> commasep (map (("#" <>) . pretty) missing) <> "."-  where-    missing =-      filter (`notElem` M.keys cs1) (M.keys cs2)-        ++ filter (`notElem` M.keys cs2) (M.keys cs1)---- | Is the given type variable the name of an abstract type or type--- parameter, which we cannot substitute?-isRigid :: VName -> Constraints -> Bool-isRigid v constraints =-  maybe True (rigidConstraint . snd) $ M.lookup v constraints---- | If the given type variable is nonrigid, what is its level?-isNonRigid :: VName -> Constraints -> Maybe Level-isNonRigid v constraints = do-  (lvl, c) <- M.lookup v constraints-  guard $ not $ rigidConstraint c-  pure lvl--type UnifySizes m =-  BreadCrumbs -> [VName] -> (VName -> Maybe Int) -> Exp -> Exp -> m ()--flipUnifySizes :: UnifySizes m -> UnifySizes m-flipUnifySizes onDims bcs bound nonrigid t1 t2 =-  onDims bcs bound nonrigid t2 t1--unifyWith ::-  (MonadUnify m) =>-  UnifySizes m ->-  Usage ->-  [VName] ->-  BreadCrumbs ->-  StructType ->-  StructType ->-  m ()-unifyWith onDims usage = subunify False-  where-    swap True x y = (y, x)-    swap False x y = (x, y)--    subunify ord bound bcs t1 t2 = do-      constraints <- getConstraints--      t1' <- normType t1-      t2' <- normType t2--      let nonrigid v = isNonRigid v constraints--          failure = matchError (srclocOf usage) mempty bcs t1' t2'--          link ord' =-            linkVarToType linkDims usage bound bcs-            where-              -- We may have to flip the order of future calls to-              -- onDims inside linkVarToType.-              linkDims-                | ord' = flipUnifySizes onDims-                | otherwise = onDims--          unifyTypeArg bcs' (TypeArgDim d1) (TypeArgDim d2) =-            onDims' bcs' (swap ord d1 d2)-          unifyTypeArg bcs' (TypeArgType t) (TypeArgType arg_t) =-            subunify ord bound bcs' t arg_t-          unifyTypeArg bcs' _ _ =-            unifyError-              usage-              mempty-              bcs'-              "Cannot unify a type argument with a dimension argument (or vice versa)."--          onDims' bcs' (d1, d2) =-            onDims-              bcs'-              bound-              nonrigid-              (applySubst (`lookupSubst` constraints) d1)-              (applySubst (`lookupSubst` constraints) d2)--      case (t1', t2') of-        (Scalar (Prim pt1), Scalar (Prim pt2))-          | pt1 == pt2 -> pure ()-        ( Scalar (Record fs),-          Scalar (Record arg_fs)-          )-            | M.keys fs == M.keys arg_fs ->-                unifySharedFields onDims usage bound bcs fs arg_fs-            | otherwise -> do-                let missing =-                      filter (`notElem` M.keys arg_fs) (M.keys fs)-                        ++ filter (`notElem` M.keys fs) (M.keys arg_fs)-                unifyError usage mempty bcs $-                  "Unshared fields:" <+> commasep (map pretty missing) <> "."-        ( Scalar (TypeVar _ (QualName _ tn) targs),-          Scalar (TypeVar _ (QualName _ arg_tn) arg_targs)-          )-            | tn == arg_tn,-              length targs == length arg_targs -> do-                let bcs' = breadCrumb (Matching "When matching type arguments.") bcs-                zipWithM_ (unifyTypeArg bcs') targs arg_targs-        ( Scalar (TypeVar _ (QualName [] v1) []),-          Scalar (TypeVar _ (QualName [] v2) [])-          ) ->-            case (nonrigid v1, nonrigid v2) of-              (Nothing, Nothing) -> failure-              (Just lvl1, Nothing) -> link ord v1 lvl1 t2'-              (Nothing, Just lvl2) -> link (not ord) v2 lvl2 t1'-              (Just lvl1, Just lvl2)-                | lvl1 <= lvl2 -> link ord v1 lvl1 t2'-                | otherwise -> link (not ord) v2 lvl2 t1'-        (Scalar (TypeVar _ (QualName [] v1) []), _)-          | Just lvl <- nonrigid v1 ->-              link ord v1 lvl t2'-        (_, Scalar (TypeVar _ (QualName [] v2) []))-          | Just lvl <- nonrigid v2 ->-              link (not ord) v2 lvl t1'-        ( Scalar (Arrow _ p1 d1 a1 (RetType b1_dims b1)),-          Scalar (Arrow _ p2 d2 a2 (RetType b2_dims b2))-          )-            | uncurry (<) $ swap ord d1 d2 -> do-                unifyError usage mempty bcs . withIndexLink "unify-consuming-param" $-                  "Parameters"-                    </> indent 2 (pretty d1 <> pretty a1)-                    </> "and"-                    </> indent 2 (pretty d2 <> pretty a2)-                    </> "are incompatible regarding consuming their arguments."-            | uncurry (<) $ swap ord (uniqueness b2) (uniqueness b1) -> do-                unifyError usage mempty bcs . withIndexLink "unify-return-uniqueness" $-                  "Return types"-                    </> indent 2 (pretty b1)-                    </> "and"-                    </> indent 2 (pretty b2)-                    </> "have incompatible uniqueness."-            | otherwise -> do-                -- Introduce the existentials as size variables so they-                -- are subject to unification.  We will remove them again-                -- afterwards.-                let (r1, r2) =-                      swap-                        ord-                        (Size Nothing $ Usage Nothing mempty)-                        (UnknownSize mempty RigidUnify)-                lvl <- curLevel-                modifyConstraints (M.fromList (map (,(lvl, r1)) b1_dims) <>)-                modifyConstraints (M.fromList (map (,(lvl, r2)) b2_dims) <>)--                let bound' = bound <> mapMaybe pname [p1, p2] <> b1_dims <> b2_dims-                subunify-                  (not ord)-                  bound-                  (breadCrumb (Matching "When matching parameter types.") bcs)-                  a1-                  a2-                subunify-                  ord-                  bound'-                  (breadCrumb (Matching "When matching return types.") bcs)-                  (toStruct b1')-                  (toStruct b2')--                -- Delete the size variables we introduced to represent-                -- the existential sizes.-                modifyConstraints $ \m -> L.foldl' (flip M.delete) m (b1_dims <> b2_dims)-            where-              (b1', b2') =-                -- Replace one parameter name with the other in the-                -- return type, in case of dependent types.  I.e.,-                -- we want type '(n: i32) -> [n]i32' to unify with-                -- type '(x: i32) -> [x]i32'.-                case (p1, p2) of-                  (Named p1', Named p2') ->-                    let f v-                          | v == p2' = Just $ ExpSubst $ sizeFromName (qualName p1') mempty-                          | otherwise = Nothing-                     in (b1, applySubst f b2)-                  (_, _) ->-                    (b1, b2)--              pname (Named x) = Just x-              pname Unnamed = Nothing-        (Array {}, Array {})-          | Shape (t1_d : _) <- arrayShape t1',-            Shape (t2_d : _) <- arrayShape t2',-            Just t1'' <- peelArray 1 t1',-            Just t2'' <- peelArray 1 t2' -> do-              onDims' bcs (swap ord t1_d t2_d)-              subunify ord bound bcs t1'' t2''-        ( Scalar (Sum cs),-          Scalar (Sum arg_cs)-          )-            | M.keys cs == M.keys arg_cs ->-                unifySharedConstructors onDims usage bound bcs cs arg_cs-            | otherwise ->-                unifyError usage mempty bcs $ unsharedConstructorsMsg arg_cs cs-        _ -> failure--anyBound :: [VName] -> ExpBase Info VName -> Bool-anyBound bound e = any (`S.member` fvVars (freeInExp e)) bound--unifySizes :: (MonadUnify m) => Usage -> UnifySizes m-unifySizes usage bcs bound nonrigid e1 e2-  | Just es <- similarExps e1 e2 =-      mapM_ (uncurry $ unifySizes usage bcs bound nonrigid) es-unifySizes usage bcs bound nonrigid (Var v1 _ _) e2-  | Just lvl1 <- nonrigid (qualLeaf v1),-    not (anyBound bound e2) || (qualLeaf v1 `elem` bound) =-      linkVarToDim usage bcs (qualLeaf v1) lvl1 e2-unifySizes usage bcs bound nonrigid e1 (Var v2 _ _)-  | Just lvl2 <- nonrigid (qualLeaf v2),-    not (anyBound bound e1) || (qualLeaf v2 `elem` bound) =-      linkVarToDim usage bcs (qualLeaf v2) lvl2 e1-unifySizes usage bcs _ _ e1 e2 = do-  notes <- (<>) <$> dimNotes usage e1 <*> dimNotes usage e2-  unifyError usage notes bcs $-    "Sizes"-      <+> dquotes (pretty e1)-      <+> "and"-      <+> dquotes (pretty e2)-      <+> "do not match."---- | Unifies two types.-unify :: (MonadUnify m) => Usage -> StructType -> StructType -> m ()-unify usage = unifyWith (unifySizes usage) usage mempty noBreadCrumbs--occursCheck ::-  (MonadUnify m) =>-  Usage ->-  BreadCrumbs ->-  VName ->-  StructType ->-  m ()-occursCheck usage bcs vn tp =-  when (vn `S.member` typeVars tp) $-    unifyError usage mempty bcs $-      "Occurs check: cannot instantiate"-        <+> prettyName vn-        <+> "with"-        <+> pretty tp-        <> "."--scopeCheck ::-  (MonadUnify m) =>-  Usage ->-  BreadCrumbs ->-  VName ->-  Level ->-  StructType ->-  m ()-scopeCheck usage bcs vn max_lvl tp = do-  constraints <- getConstraints-  checkType constraints tp-  where-    checkType constraints t =-      mapM_ (check constraints) $ typeVars t <> fvVars (freeInType t)--    check constraints v-      | Just (lvl, c) <- M.lookup v constraints,-        lvl > max_lvl =-          if rigidConstraint c-            then scopeViolation v-            else modifyConstraints $ M.insert v (max_lvl, c)-      | otherwise =-          pure ()--    scopeViolation v = do-      notes <- typeNotes usage tp-      unifyError usage notes bcs $-        "Cannot unify type"-          </> indent 2 (pretty tp)-          </> "with"-          <+> dquotes (prettyName vn)-          <+> "(scope violation)."-          </> "This is because"-          <+> dquotes (prettyName v)-          <+> "is rigidly bound in a deeper scope."---- Expressions witnessed by type, topologically sorted.-topWit :: TypeBase Exp u -> [Exp]-topWit = topologicalSort depends . witnessedExps-  where-    witnessedExps t = execState (traverseDims onDim t) mempty-      where-        onDim _ PosImmediate e = modify (e :)-        onDim _ _ _ = pure ()-    depends a b = any (sameExp b) $ subExps a--sizeFree ::-  (MonadUnify m) =>-  SrcLoc ->-  (Exp -> Maybe VName) ->-  TypeBase Size u ->-  m (TypeBase Size u, [VName])-sizeFree tloc expKiller orig_t = do-  runReaderT (toBeReplaced orig_t $ onType orig_t) mempty `runStateT` mempty-  where-    lookReplacement e repl = snd <$> L.find (sameExp e . fst) repl-    expReplace mapping e-      | Just e' <- lookReplacement e mapping = e'-      | otherwise = runIdentity $ astMap mapper e-      where-        mapper = identityMapper {mapOnExp = pure . expReplace mapping}--    replacing e = do-      e' <- asks (`expReplace` e)-      case expKiller e' of-        Nothing -> pure e'-        Just cause -> do-          vn <- lift $ lift $ newRigidDim tloc (RigidOutOfScope (locOf e) cause) "d"-          modify (vn :)-          pure $ sizeFromName (qualName vn) (srclocOf e)--    toBeReplaced t m' = foldl f m' $ topWit t-      where-        f m e = do-          e' <- replacing e-          local ((e, e') :) m--    onScalar (Record fs) =-      Record <$> traverse onType fs-    onScalar (Sum cs) =-      Sum <$> (traverse . traverse) onType cs-    onScalar (Arrow as pn d argT (RetType dims retT)) = do-      argT' <- onType argT-      old_bound <- get-      retT' <- toBeReplaced retT $ onType retT-      rl <- state $ L.partition (`notElem` old_bound)-      let dims' = dims <> rl-      pure $ Arrow as pn d argT' (RetType dims' retT')-    onScalar (TypeVar u v args) =-      TypeVar u v <$> mapM onTypeArg args-      where-        onTypeArg (TypeArgDim d) = TypeArgDim <$> replacing d-        onTypeArg (TypeArgType ty) = TypeArgType <$> onType ty-    onScalar (Prim pt) = pure $ Prim pt--    onType ::-      (MonadUnify m) =>-      TypeBase Size u ->-      ReaderT [(Exp, Exp)] (StateT [VName] m) (TypeBase Size u)-    onType (Array u shape scalar) =-      Array u <$> traverse replacing shape <*> onScalar scalar-    onType (Scalar ty) =-      Scalar <$> onScalar ty--linkVarToType ::-  (MonadUnify m) =>-  UnifySizes m ->-  Usage ->-  [VName] ->-  BreadCrumbs ->-  VName ->-  Level ->-  StructType ->-  m ()-linkVarToType onDims usage bound bcs vn lvl tp_unnorm = do-  -- We have to expand anyway for the occurs check, so we might as-  -- well link the fully expanded type.-  tp <- normTypeFully tp_unnorm-  occursCheck usage bcs vn tp-  scopeCheck usage bcs vn lvl tp--  let link = do-        let (witnessed, not_witnessed) = determineSizeWitnesses tp-            used v = v `S.member` witnessed || v `S.member` not_witnessed-            (ext_witnessed, ext_not_witnessed) =-              L.partition (`elem` witnessed) $ filter used bound--            -- Any size that uses an ext_not_witnessed variable must-            -- be replaced with a fresh existential.-            problematic e =-              L.find (`elem` ext_not_witnessed) $-                S.toList $-                  fvVars $-                    freeInExp e--        (tp', ext_new) <- sizeFree (srclocOf usage) problematic tp--        modifyConstraints $-          M.insert vn (lvl, Constraint (RetType (ext_new <> ext_witnessed) tp') usage)--  let unliftedBcs unlifted_usage =-        breadCrumb-          ( Matching $-              "When verifying that"-                <+> dquotes (prettyName vn)-                <+> textwrap "is not instantiated with a function type, due to"-                <+> pretty unlifted_usage-          )-          bcs--  constraints <- getConstraints-  case snd <$> M.lookup vn constraints of-    Just (NoConstraint Unlifted unlift_usage) -> do-      link--      arrayElemTypeWith usage (unliftedBcs unlift_usage) tp-      when (any (`elem` bound) (fvVars (freeInType tp))) $-        unifyError usage mempty bcs $-          "Type variable"-            <+> prettyName vn-            <+> "cannot be instantiated with type containing anonymous sizes:"-            </> indent 2 (pretty tp)-            </> textwrap "This is usually because the size of an array returned by a higher-order function argument cannot be determined statically.  This can also be due to the return size being a value parameter.  Add type annotation to clarify."-    Just (Equality _) -> do-      link-      equalityType usage tp-    Just (Overloaded ts old_usage)-      | tp `notElem` map (Scalar . Prim) ts -> do-          link-          case tp of-            Scalar (TypeVar _ (QualName [] v) [])-              | not $ isRigid v constraints ->-                  linkVarToTypes usage v ts-            _ ->-              unifyError usage mempty bcs $-                "Cannot instantiate"-                  <+> dquotes (prettyName vn)-                  <+> "with type"-                  </> indent 2 (pretty tp)-                  </> "as"-                  <+> dquotes (prettyName vn)-                  <+> "must be one of"-                  <+> commasep (map pretty ts)-                  </> "due to"-                  <+> pretty old_usage-                  <> "."-    Just (HasFields l required_fields old_usage) -> do-      when (l == Unlifted) $ arrayElemTypeWith usage (unliftedBcs old_usage) tp-      case tp of-        Scalar (Record tp_fields)-          | all (`M.member` tp_fields) $ M.keys required_fields -> do-              required_fields' <- mapM normTypeFully required_fields-              let tp' = Scalar $ Record $ required_fields <> tp_fields -- Crucially left-biased.-                  ext = filter (`S.member` fvVars (freeInType tp')) bound-              modifyConstraints $-                M.insert vn (lvl, Constraint (RetType ext tp') usage)-              unifySharedFields onDims usage bound bcs required_fields' tp_fields-        Scalar (TypeVar _ (QualName [] v) []) -> do-          case M.lookup v constraints of-            Just (_, HasFields _ tp_fields _) ->-              unifySharedFields onDims usage bound bcs required_fields tp_fields-            Just (_, NoConstraint {}) -> pure ()-            Just (_, Equality {}) -> pure ()-            _ -> do-              notes <- (<>) <$> typeVarNotes vn <*> typeVarNotes v-              noRecordType notes-          link-          modifyConstraints $-            M.insertWith-              combineFields-              v-              (lvl, HasFields l required_fields old_usage)-          where-            combineFields (_, HasFields l1 fs1 usage1) (_, HasFields l2 fs2 _) =-              (lvl, HasFields (l1 `min` l2) (M.union fs1 fs2) usage1)-            combineFields hasfs _ = hasfs-        _ ->-          unifyError usage mempty bcs $-            "Cannot instantiate"-              <+> dquotes (prettyName vn)-              <+> "with type"-              </> indent 2 (pretty tp)-              </> "as"-              <+> dquotes (prettyName vn)-              <+> "must be a record with fields"-              </> indent 2 (pretty (Record required_fields))-              </> "due to"-              <+> pretty old_usage-              <> "."-    -- See Note [Linking variables to sum types]-    Just (HasConstrs l required_cs old_usage) -> do-      when (l == Unlifted) $ arrayElemTypeWith usage (unliftedBcs old_usage) tp-      case tp of-        Scalar (Sum ts)-          | all (`M.member` ts) $ M.keys required_cs -> do-              let tp' = Scalar $ Sum $ required_cs <> ts -- Crucially left-biased.-                  ext = filter (`S.member` fvVars (freeInType tp')) bound-              modifyConstraints $-                M.insert vn (lvl, Constraint (RetType ext tp') usage)-              unifySharedConstructors onDims usage bound bcs required_cs ts-          | otherwise ->-              unsharedConstructors required_cs ts =<< typeVarNotes vn-        Scalar (TypeVar _ (QualName [] v) []) -> do-          case M.lookup v constraints of-            Just (_, HasConstrs _ v_cs _) ->-              unifySharedConstructors onDims usage bound bcs required_cs v_cs-            Just (_, NoConstraint {}) -> pure ()-            Just (_, Equality {}) -> pure ()-            _ -> do-              notes <- (<>) <$> typeVarNotes vn <*> typeVarNotes v-              noSumType notes-          link-          modifyConstraints $-            M.insertWith-              combineConstrs-              v-              (lvl, HasConstrs l required_cs old_usage)-          where-            combineConstrs (_, HasConstrs l1 cs1 usage1) (_, HasConstrs l2 cs2 _) =-              (lvl, HasConstrs (l1 `min` l2) (M.union cs1 cs2) usage1)-            combineConstrs hasCs _ = hasCs-        _ -> noSumType =<< typeVarNotes vn-    _ -> link-  where-    unsharedConstructors cs1 cs2 notes =-      unifyError-        usage-        notes-        bcs-        (unsharedConstructorsMsg cs1 cs2)-    noSumType notes =-      unifyError-        usage-        notes-        bcs-        "Cannot unify a sum type with a non-sum type."-    noRecordType notes =-      unifyError-        usage-        notes-        bcs-        "Cannot unify a record type with a non-record type."--linkVarToDim ::-  (MonadUnify m) =>-  Usage ->-  BreadCrumbs ->-  VName ->-  Level ->-  Exp ->-  m ()-linkVarToDim usage bcs vn lvl e = do-  constraints <- getConstraints--  mapM_ (checkVar constraints) $ fvVars $ freeInExp e--  modifyConstraints $ M.insert vn (lvl, Size (Just e) usage)-  where-    checkVar _ dim'-      | vn == dim' = do-          notes <- dimNotes usage e-          unifyError usage notes bcs $-            "Occurs check: cannot instantiate"-              <+> dquotes (prettyName vn)-              <+> "with"-              <+> dquotes (pretty e)-              <+> "."-    checkVar constraints dim'-      | Just (dim_lvl, c) <- dim' `M.lookup` constraints,-        dim_lvl >= lvl =-          case c of-            ParamSize {} -> do-              notes <- dimNotes usage e-              unifyError usage notes bcs $-                "Cannot link size"-                  <+> dquotes (prettyName vn)-                  <+> "to"-                  <+> dquotes (pretty e)-                  <+> "(scope violation)."-                  </> "This is because"-                  <+> dquotes (pretty $ qualName dim')-                  <+> "is not in scope when"-                  <+> dquotes (prettyName vn)-                  <+> "is introduced."-            _ -> modifyConstraints $ M.insert dim' (lvl, c)-    checkVar _ _ = pure ()---- | Assert that this type must be one of the given primitive types.-mustBeOneOf :: (MonadUnify m) => [PrimType] -> Usage -> StructType -> m ()-mustBeOneOf [req_t] usage t = unify usage (Scalar (Prim req_t)) t-mustBeOneOf ts usage t = do-  t' <- normType t-  constraints <- getConstraints-  let isRigid' v = isRigid v constraints--  case t' of-    Scalar (TypeVar _ (QualName [] v) [])-      | not $ isRigid' v -> linkVarToTypes usage v ts-    Scalar (Prim pt) | pt `elem` ts -> pure ()-    _ -> failure-  where-    failure =-      unifyError usage mempty noBreadCrumbs $-        "Cannot unify type"-          <+> dquotes (pretty t)-          <+> "with any of "-          <> commasep (map pretty ts)-          <> "."--linkVarToTypes :: (MonadUnify m) => Usage -> VName -> [PrimType] -> m ()-linkVarToTypes usage vn ts = do-  vn_constraint <- M.lookup vn <$> getConstraints-  case vn_constraint of-    Just (lvl, Overloaded vn_ts vn_usage) ->-      case ts `L.intersect` vn_ts of-        [] ->-          unifyError usage mempty noBreadCrumbs $-            "Type constrained to one of"-              <+> commasep (map pretty ts)-              <+> "but also one of"-              <+> commasep (map pretty vn_ts)-              <+> "due to"-              <+> pretty vn_usage-              <> "."-        ts' -> modifyConstraints $ M.insert vn (lvl, Overloaded ts' usage)-    Just (_, HasConstrs _ _ vn_usage) ->-      unifyError usage mempty noBreadCrumbs $-        "Type constrained to one of"-          <+> commasep (map pretty ts)-          <> ", but also inferred to be sum type due to"-            <+> pretty vn_usage-          <> "."-    Just (_, HasFields _ _ vn_usage) ->-      unifyError usage mempty noBreadCrumbs $-        "Type constrained to one of"-          <+> commasep (map pretty ts)-          <> ", but also inferred to be record due to"-            <+> pretty vn_usage-          <> "."-    Just (lvl, _) -> modifyConstraints $ M.insert vn (lvl, Overloaded ts usage)-    Nothing ->-      unifyError usage mempty noBreadCrumbs $-        "Cannot constrain type to one of" <+> commasep (map pretty ts)---- | Assert that this type must support equality.-equalityType ::-  (MonadUnify m, Pretty (Shape dim), Pretty u) =>-  Usage ->-  TypeBase dim u ->-  m ()-equalityType usage t = do-  unless (orderZero t) $-    unifyError usage mempty noBreadCrumbs $-      "Type " <+> dquotes (pretty t) <+> "does not support equality (may contain function)."-  mapM_ mustBeEquality $ typeVars t-  where-    mustBeEquality vn = do-      constraints <- getConstraints-      case M.lookup vn constraints of-        Just (_, Constraint (RetType [] (Scalar (TypeVar _ (QualName [] vn') []))) _) ->-          mustBeEquality vn'-        Just (_, Constraint (RetType _ vn_t) cusage)-          | not $ orderZero vn_t ->-              unifyError usage mempty noBreadCrumbs $-                "Type"-                  <+> dquotes (pretty t)-                  <+> "does not support equality."-                  </> "Constrained to be higher-order due to"-                  <+> pretty cusage-                  <+> "."-          | otherwise -> pure ()-        Just (lvl, NoConstraint _ _) ->-          modifyConstraints $ M.insert vn (lvl, Equality usage)-        Just (_, Overloaded _ _) ->-          pure () -- All primtypes support equality.-        Just (_, Equality {}) ->-          pure ()-        _ ->-          unifyError usage mempty noBreadCrumbs $-            "Type" <+> prettyName vn <+> "does not support equality."--zeroOrderTypeWith ::-  (MonadUnify m) =>-  Usage ->-  BreadCrumbs ->-  StructType ->-  m ()-zeroOrderTypeWith usage bcs t = do-  unless (orderZero t) $-    unifyError usage mempty bcs $-      "Type" </> indent 2 (pretty t) </> "found to be functional."-  mapM_ mustBeZeroOrder . S.toList . typeVars =<< normType t-  where-    mustBeZeroOrder vn = do-      constraints <- getConstraints-      case M.lookup vn constraints of-        Just (lvl, NoConstraint _ _) ->-          modifyConstraints $ M.insert vn (lvl, NoConstraint Unlifted usage)-        Just (lvl, HasFields _ fs _) ->-          modifyConstraints $ M.insert vn (lvl, HasFields Unlifted fs usage)-        Just (lvl, HasConstrs _ cs _) ->-          modifyConstraints $ M.insert vn (lvl, HasConstrs Unlifted cs usage)-        Just (_, ParamType Lifted ploc) ->-          unifyError usage mempty bcs $-            "Type parameter"-              <+> dquotes (prettyName vn)-              <+> "at"-              <+> pretty (locStr ploc)-              <+> "may be a function."-        _ -> pure ()---- | Assert that this type must be zero-order.-zeroOrderType ::-  (MonadUnify m) => Usage -> T.Text -> StructType -> m ()-zeroOrderType usage desc =-  zeroOrderTypeWith usage $ breadCrumb bc noBreadCrumbs-  where-    bc = Matching $ "When checking" <+> textwrap desc--arrayElemTypeWith ::-  (MonadUnify m, Pretty (Shape dim), Pretty u) =>-  Usage ->-  BreadCrumbs ->-  TypeBase dim u ->-  m ()-arrayElemTypeWith usage bcs t = do-  unless (orderZero t) $-    unifyError usage mempty bcs $-      "Type" </> indent 2 (pretty t) </> "found to be functional."-  mapM_ mustBeZeroOrder . S.toList . typeVars $ t-  where-    mustBeZeroOrder vn = do-      constraints <- getConstraints-      case M.lookup vn constraints of-        Just (lvl, NoConstraint _ _) ->-          modifyConstraints $ M.insert vn (lvl, NoConstraint Unlifted usage)-        Just (_, ParamType l ploc)-          | l `elem` [Lifted, SizeLifted] ->-              unifyError usage mempty bcs $-                "Type parameter"-                  <+> dquotes (prettyName vn)-                  <+> "bound at"-                  <+> pretty (locStr ploc)-                  <+> "is lifted and cannot be an array element."-        _ -> pure ()---- | Assert that this type must be valid as an array element.-arrayElemType ::-  (MonadUnify m, Pretty (Shape dim), Pretty u) =>-  Usage ->-  T.Text ->-  TypeBase dim u ->-  m ()-arrayElemType usage desc =-  arrayElemTypeWith usage $ breadCrumb bc noBreadCrumbs-  where-    bc = Matching $ "When checking" <+> textwrap desc--unifySharedFields ::-  (MonadUnify m) =>-  UnifySizes m ->-  Usage ->-  [VName] ->-  BreadCrumbs ->-  M.Map Name StructType ->-  M.Map Name StructType ->-  m ()-unifySharedFields onDims usage bound bcs fs1 fs2 =-  forM_ (M.toList $ M.intersectionWith (,) fs1 fs2) $ \(f, (t1, t2)) ->-    unifyWith onDims usage bound (breadCrumb (MatchingFields [f]) bcs) t1 t2--unifySharedConstructors ::-  (MonadUnify m) =>-  UnifySizes m ->-  Usage ->-  [VName] ->-  BreadCrumbs ->-  M.Map Name [StructType] ->-  M.Map Name [StructType] ->-  m ()-unifySharedConstructors onDims usage bound bcs cs1 cs2 =-  forM_ (M.toList $ M.intersectionWith (,) cs1 cs2) $ \(c, (f1, f2)) ->-    unifyConstructor c f1 f2-  where-    unifyConstructor c f1 f2-      | length f1 == length f2 = do-          let bcs' = breadCrumb (MatchingConstructor c) bcs-          zipWithM_ (unifyWith onDims usage bound bcs') f1 f2-      | otherwise =-          unifyError usage mempty bcs $-            "Cannot unify constructor" <+> dquotes (prettyName c) <> "."---- | In @mustHaveConstr usage c t fs@, the type @t@ must have a--- constructor named @c@ that takes arguments of types @ts@.-mustHaveConstr ::-  (MonadUnify m) =>-  Usage ->-  Name ->-  StructType ->-  [StructType] ->-  m ()-mustHaveConstr usage c t fs = do-  constraints <- getConstraints-  case t of-    Scalar (TypeVar _ (QualName _ tn) [])-      | Just (lvl, NoConstraint l _) <- M.lookup tn constraints -> do-          mapM_ (scopeCheck usage noBreadCrumbs tn lvl) fs-          modifyConstraints $ M.insert tn (lvl, HasConstrs l (M.singleton c fs) usage)-      | Just (lvl, HasConstrs l cs _) <- M.lookup tn constraints ->-          case M.lookup c cs of-            Nothing ->-              modifyConstraints $-                M.insert tn (lvl, HasConstrs l (M.insert c fs cs) usage)-            Just fs'-              | length fs == length fs' -> zipWithM_ (unify usage) fs fs'-              | otherwise ->-                  unifyError usage mempty noBreadCrumbs $-                    "Different arity for constructor" <+> dquotes (pretty c) <> "."-    Scalar (Sum cs) ->-      case M.lookup c cs of-        Nothing ->-          unifyError usage mempty noBreadCrumbs $-            "Constuctor" <+> dquotes (pretty c) <+> "not present in type."-        Just fs'-          | length fs == length fs' -> zipWithM_ (unify usage) fs fs'-          | otherwise ->-              unifyError usage mempty noBreadCrumbs $-                "Different arity for constructor" <+> dquotes (pretty c) <+> "."-    _ ->-      unify usage t $ Scalar $ Sum $ M.singleton c fs--mustHaveFieldWith ::-  (MonadUnify m) =>-  UnifySizes m ->-  Usage ->-  [VName] ->-  BreadCrumbs ->-  Name ->-  StructType ->-  m StructType-mustHaveFieldWith onDims usage bound bcs l t = do-  constraints <- getConstraints-  l_type <- newTypeVar (locOf usage) "t"-  case t of-    Scalar (TypeVar _ (QualName _ tn) [])-      | Just (lvl, NoConstraint {}) <- M.lookup tn constraints -> do-          scopeCheck usage bcs tn lvl l_type-          modifyConstraints $ M.insert tn (lvl, HasFields Lifted (M.singleton l l_type) usage)-          pure l_type-      | Just (lvl, HasFields lifted fields _) <- M.lookup tn constraints -> do-          case M.lookup l fields of-            Just t' -> unifyWith onDims usage bound bcs l_type t'-            Nothing ->-              modifyConstraints $-                M.insert-                  tn-                  (lvl, HasFields lifted (M.insert l l_type fields) usage)-          pure l_type-    Scalar (Record fields)-      | Just t' <- M.lookup l fields -> do-          unify usage l_type t'-          pure t'-      | otherwise ->-          unifyError usage mempty bcs $-            "Attempt to access field"-              <+> dquotes (pretty l)-              <+> " of value of type"-              <+> pretty (toStructural t)-              <> "."-    _ -> do-      unify usage t $ Scalar $ Record $ M.singleton l l_type-      pure l_type---- | Assert that some type must have a field with this name and type.-mustHaveField ::-  (MonadUnify m) =>-  Usage ->-  Name ->-  StructType ->-  m StructType-mustHaveField usage = mustHaveFieldWith (unifySizes usage) usage mempty noBreadCrumbs--newDimOnMismatch ::-  (MonadUnify m) =>-  Loc ->-  StructType ->-  StructType ->-  m (StructType, [VName])-newDimOnMismatch loc t1 t2 = do-  (t, seen) <- runStateT (matchDims onDims t1 t2) mempty-  pure (t, M.elems seen)-  where-    r = RigidCond t1 t2-    same (e1, e2) =-      maybe False (all same) $ similarExps e1 e2-    onDims _ d1 d2-      | same (d1, d2) = pure d1-      | otherwise = do-          -- Remember mismatches we have seen before and reuse the-          -- same new size.-          maybe_d <- gets $ M.lookup (d1, d2)-          case maybe_d of-            Just d -> pure $ sizeFromName (qualName d) $ srclocOf loc-            Nothing -> do-              d <- lift $ newRigidDim loc r "differ"-              modify $ M.insert (d1, d2) d-              pure $ sizeFromName (qualName d) $ srclocOf loc---- | Like unification, but creates new size variables where mismatches--- occur.  Returns the new dimensions thus created.-unifyMostCommon ::-  (MonadUnify m) =>-  Usage ->-  StructType ->-  StructType ->-  m (StructType, [VName])-unifyMostCommon usage t1 t2 = do-  -- We are ignoring the dimensions here, because any mismatches-  -- should be turned into fresh size variables.-  let allOK _ _ _ _ _ = pure ()-  unifyWith allOK usage mempty noBreadCrumbs t1 t2-  t1' <- normTypeFully t1-  t2' <- normTypeFully t2-  newDimOnMismatch (locOf usage) t1' t2'---- Simple MonadUnify implementation.--type UnifyMState = (Constraints, Int)--newtype UnifyM a = UnifyM (StateT UnifyMState (Except TypeError) a)-  deriving-    ( Monad,-      Functor,-      Applicative,-      MonadState UnifyMState,-      MonadError TypeError-    )--newVar :: Name -> UnifyM VName-newVar name = do-  (x, i) <- get-  put (x, i + 1)-  pure $ VName (mkTypeVarName name i) i--instance MonadUnify UnifyM where-  getConstraints = gets fst-  putConstraints x = modify $ \(_, i) -> (x, i)--  newTypeVar loc name = do-    v <- newVar name-    modifyConstraints $ M.insert v (0, NoConstraint Lifted $ Usage Nothing $ locOf loc)-    pure $ Scalar $ TypeVar mempty (qualName v) []--  newDimVar usage rigidity name = do-    dim <- newVar name-    case rigidity of-      Rigid src ->-        modifyConstraints $-          M.insert dim (0, UnknownSize (locOf usage) src)-      Nonrigid ->-        modifyConstraints $-          M.insert dim (0, Size Nothing usage)-    pure dim--  curLevel = pure 1--  unifyError loc notes bcs doc =-    throwError $ TypeError (locOf loc) notes $ doc <> pretty bcs--  matchError loc notes bcs t1 t2 =-    throwError $ TypeError (locOf loc) notes $ doc <> pretty bcs-    where-      doc =-        "Types"-          </> indent 2 (pretty t1)-          </> "and"-          </> indent 2 (pretty t2)-          </> "do not match."--runUnifyM :: [TypeParam] -> [TypeParam] -> UnifyM a -> Either TypeError a-runUnifyM rigid_tparams nonrigid_tparams (UnifyM m) =-  runExcept $ evalStateT m (constraints, 0)-  where-    constraints =-      M.fromList $-        map nonrigid nonrigid_tparams <> map rigid rigid_tparams-    nonrigid (TypeParamDim p loc) = (p, (1, Size Nothing $ Usage Nothing $ locOf loc))-    nonrigid (TypeParamType l p loc) = (p, (1, NoConstraint l $ Usage Nothing $ locOf loc))-    rigid (TypeParamDim p loc) = (p, (0, ParamSize $ locOf loc))-    rigid (TypeParamType l p loc) = (p, (0, ParamType l $ locOf loc))---- | Perform a unification of two types outside a monadic context.--- The first list of type parameters are rigid but may have liftedness--- constraints; the second list of type parameters are allowed to be--- instantiated. All other types are considered rigid with no--- constraints.-doUnification ::-  Loc ->-  [TypeParam] ->-  [TypeParam] ->-  StructType ->-  StructType ->-  Either TypeError StructType-doUnification loc rigid_tparams nonrigid_tparams t1 t2 =-  runUnifyM rigid_tparams nonrigid_tparams $ do-    unify (Usage Nothing (locOf loc)) t1 t2-    normTypeFully t2---- Note [Linking variables to sum types]------ Consider the case when unifying a result type------   i32 -> ?[n].(#foo [n]bool)------ with------   i32 -> ?[k].a------ where 'a' has a HasConstrs constraint saying that it must have at--- least a constructor of type '#foo [0]bool'.------ This unification should succeed, but we must not merely link 'a' to--- '#foo [n]bool', as 'n' is not free.  Instead we should instantiate--- 'a' to be a concrete sum type (because now we know exactly which--- constructor labels it must have), and unify each of its constructor--- payloads with the corresponding expected payload.+{-# LANGUAGE LambdaCase #-}++-- | Implementation of unification and other core type system building+-- blocks.+module Language.Futhark.TypeChecker.Unify+  ( Constraint (..),+    Usage (..),+    mkUsage,+    mkUsage',+    Level,+    Constraints,+    MonadUnify (..),+    Rigidity (..),+    RigidSource (..),+    BreadCrumbs,+    allDimsFreshInType,+    dimNotes,+    normTypeFully,+    unify,+    unifyMostCommon,+    doUnification,+  )+where++import Control.Monad+import Control.Monad.Except+import Control.Monad.Reader+import Control.Monad.State+import Data.Bifunctor+import Data.Bitraversable+import Data.List qualified as L+import Data.Map.Strict qualified as M+import Data.Maybe+import Data.Set qualified as S+import Data.Text qualified as T+import Futhark.Util.Pretty+import Language.Futhark+import Language.Futhark.TypeChecker.Constraints (CtTy (..), Level, Reason (..), TyVarInfo (..))+import Language.Futhark.TypeChecker.Error+import Language.Futhark.TypeChecker.Monad hiding (BoundV)+import Language.Futhark.TypeChecker.TySolve qualified as TySolve+import Language.Futhark.TypeChecker.Types++-- | A usage that caused a type constraint.+data Usage = Usage (Maybe T.Text) Loc+  deriving (Show)++-- | Construct a 'Usage' from a location and a description.+mkUsage :: (Located a) => a -> T.Text -> Usage+mkUsage = flip (Usage . Just) . locOf++-- | Construct a 'Usage' that has just a location, but no particular+-- description.+mkUsage' :: (Located a) => a -> Usage+mkUsage' = Usage Nothing . locOf++instance Pretty Usage where+  pretty (Usage Nothing loc) = "use at " <> textwrap (locText loc)+  pretty (Usage (Just s) loc) = textwrap s <+> "at" <+> textwrap (locText loc)++instance Located Usage where+  locOf (Usage _ loc) = locOf loc++-- | A constraint on a yet-ambiguous size variable, or information+-- about a rigid type parameter or size.+data Constraint+  = ParamSize Loc+  | -- | Is not actually a type, but a term-level size,+    -- possibly already set to something specific.+    Size (Maybe Exp) Usage+  | -- | A size that does not unify with anything -+    -- created from the result of applying a function+    -- whose return size is existential, or otherwise+    -- hiding a size.+    UnknownSize Loc RigidSource+  | -- | A size arising from instantiating a type parameter (of the+    -- given liftedness) with a type whose sizes are not yet known.+    -- In contrast to an ordinary 'Size', unification may determine+    -- that this size is actually existential (unless the type+    -- parameter is unlifted), in which case the constraint is+    -- replaced with 'ExistentialSize'. See Note [Size Inference] in+    -- Language.Futhark.TypeChecker.Terms.+    InstSize Liftedness Usage+  | -- | Another occurrence of the instantiated size denoted by the+    -- given canonical size variable (an 'InstSize'). Kept distinct+    -- from the canonical variable because if the size turns out to+    -- be existential, every occurrence must be a distinct+    -- existential. The integer identifies the occurrence of the+    -- instantiated type parameter that this size is part of; copies+    -- from the same occurrence absorbed from the same source denote+    -- the same existential size.+    CopySize VName Int Usage+  | -- | An instantiated size that unification has determined to+    -- correspond to an existential size, possibly with the variable+    -- it was unified with (used to identify existentials with the+    -- same origin). Each variable constrained by this (or a+    -- 'CopySize' pointing to it) is turned into a rigid size when+    -- the enclosing function application is complete. The second+    -- field, when present, means that a *rigid* unknown size was+    -- absorbed: the existential then stands for a size that is+    -- actually computed at that location, and if it cannot be bound+    -- anywhere, it becomes a rigid unknown size there (see+    -- bindExistentialInsts), subjecting it to the causality check.+    ExistentialSize (Maybe VName) (Maybe Loc) Usage+  deriving (Show)++instance Located Constraint where+  locOf (ParamSize loc) = locOf loc+  locOf (Size _ usage) = locOf usage+  locOf (UnknownSize loc _) = locOf loc+  locOf (InstSize _ usage) = locOf usage+  locOf (CopySize _ _ usage) = locOf usage+  locOf (ExistentialSize _ _ usage) = locOf usage++-- | Mapping from fresh type variables, instantiated from the type+-- schemes of polymorphic functions, to (possibly) specific types as+-- determined on application and the location of that application, or+-- a partial constraint on their type.+type Constraints = M.Map VName (Level, Constraint)++lookupSubst :: VName -> Constraints -> Maybe (Subst StructRetType)+lookupSubst v constraints = case snd <$> M.lookup v constraints of+  Just (Size (Just d) _) ->+    Just $ ExpSubst $ applySubst (`lookupSubst` constraints) d+  Just (CopySize c _ _)+    -- If the canonical size has been resolved to an actual size, we+    -- are equal to that size. Otherwise (canonical size still+    -- pending, or existential) we stand apart under our own name.+    | Just (Size (Just _) _) <- snd <$> M.lookup c constraints ->+        lookupSubst c constraints+  _ -> Nothing++-- | The source of a rigid size.+data RigidSource+  = -- | A function argument that is not a constant or variable name.+    RigidArg (Maybe (QualName VName)) T.Text+  | -- | An existential return size.+    RigidRet (Maybe (QualName VName))+  | -- | Similarly to 'RigidRet', but produce by a loop.+    RigidLoop+  | -- | Produced by a complicated slice expression.+    RigidSlice (Maybe Size) T.Text+  | -- | Produced by a complicated range expression.+    RigidRange+  | -- | Mismatch in branches.+    RigidCond StructType StructType+  | -- | Invented during unification.+    RigidUnify+  | -- | A name used in a size went out of scope.+    RigidOutOfScope Loc VName+  deriving (Eq, Ord, Show)++-- | The ridigity of a size variable.  All rigid sizes are tagged with+-- information about how they were generated.+data Rigidity = Rigid RigidSource | Nonrigid+  deriving (Eq, Ord, Show)++prettySource :: Loc -> Loc -> RigidSource -> Doc ()+prettySource ctx loc (RigidRet Nothing) =+  "is unknown size returned by function at"+    <+> pretty (locStrRel ctx loc)+    <> "."+prettySource ctx loc (RigidRet (Just fname)) =+  "is unknown size returned by"+    <+> dquotes (pretty fname)+    <+> "at"+    <+> pretty (locStrRel ctx loc)+    <> "."+prettySource ctx loc (RigidArg fname arg) =+  "is value of argument"+    </> indent 2 (shorten (pretty arg))+    </> "passed to"+    <+> fname'+    <+> "at"+    <+> pretty (locStrRel ctx loc)+    <> "."+  where+    fname' = maybe "function" (dquotes . pretty) fname+prettySource ctx loc (RigidSlice d slice) =+  "is size produced by slice"+    </> indent 2 (shorten (pretty slice))+    </> d_desc+    <> "at"+      <+> pretty (locStrRel ctx loc)+    <> "."+  where+    d_desc = case d of+      Just d' -> "of dimension of size " <> dquotes (pretty d') <> " "+      Nothing -> mempty+prettySource ctx loc RigidLoop =+  "is unknown size of value returned at" <+> pretty (locStrRel ctx loc) <> "."+prettySource ctx loc RigidRange =+  "is unknown length of range at" <+> pretty (locStrRel ctx loc) <> "."+prettySource ctx loc (RigidOutOfScope boundloc v) =+  "is an unknown size arising from "+    <> dquotes (prettyName v)+    <> " going out of scope at "+    <> pretty (locStrRel ctx loc)+    <> "."+      </> "Originally bound at "+    <> pretty (locStrRel ctx boundloc)+    <> "."+prettySource _ _ RigidUnify =+  textwrap "is an artificial size invented during unification of functions with anonymous sizes."+prettySource ctx loc (RigidCond t1 t2) =+  "is unknown due to conditional expression at "+    <> pretty (locStrRel ctx loc)+    <> "."+      </> "One branch returns array of type: "+    <> align (pretty t1)+      </> "The other an array of type:       "+    <> align (pretty t2)++-- | Retrieve notes describing the purpose or origin of the given+-- t'Size'.  The location is used as the *current* location, for the+-- purpose of reporting relative locations.+dimNotes :: (Located a, MonadUnify m) => a -> Exp -> m Notes+dimNotes ctx (Var d _ _) = do+  c <- M.lookup (qualLeaf d) <$> getConstraints+  case c of+    Just (_, UnknownSize loc rsrc) ->+      pure . aNote $+        dquotes (pretty d) <+> prettySource (locOf ctx) loc rsrc+    _ -> pure mempty+dimNotes _ _ = pure mempty++-- | Monads that which to perform unification must implement this type+-- class.+class (Monad m) => MonadUnify m where+  getConstraints :: m Constraints+  putConstraints :: Constraints -> m ()+  modifyConstraints :: (Constraints -> Constraints) -> m ()+  modifyConstraints f = do+    x <- getConstraints+    putConstraints $ f x++  newDimVar :: Usage -> Rigidity -> Name -> m VName+  newRigidDim :: (Located a) => a -> RigidSource -> Name -> m VName+  newRigidDim loc = newDimVar (mkUsage' loc) . Rigid+  newFlexibleDim :: Usage -> Name -> m VName+  newFlexibleDim usage = newDimVar usage Nonrigid++  curLevel :: m Level++  matchError ::+    (Located loc) =>+    loc ->+    Notes ->+    BreadCrumbs ->+    StructType ->+    StructType ->+    m a++  unifyError ::+    (Located loc) =>+    loc ->+    Notes ->+    BreadCrumbs ->+    Doc () ->+    m a++-- | Replace all type variables with their substitution.+normTypeFully :: (Substitutable a, MonadUnify m) => a -> m a+normTypeFully t = do+  constraints <- getConstraints+  pure $ applySubst (`lookupSubst` constraints) t++rigidConstraint :: Constraint -> Bool+rigidConstraint ParamSize {} = True+rigidConstraint UnknownSize {} = True+rigidConstraint ExistentialSize {} = True+rigidConstraint _ = False++-- | If the given type variable is nonrigid, what is its level?+isNonRigid :: VName -> Constraints -> Maybe Level+isNonRigid v constraints = do+  (lvl, c) <- M.lookup v constraints+  case c of+    -- A copy is as rigid as its canonical size.+    CopySize c' _ _ | Just (_, c'') <- M.lookup c' constraints -> do+      guard $ not $ rigidConstraint c''+      pure lvl+    _ -> do+      guard $ not $ rigidConstraint c+      pure lvl++type UnifySizes m =+  BreadCrumbs -> [VName] -> (VName -> Maybe Int) -> Exp -> Exp -> m ()++unifyWith ::+  (MonadUnify m) =>+  UnifySizes m ->+  Usage ->+  [VName] ->+  BreadCrumbs ->+  StructType ->+  StructType ->+  m ()+unifyWith onDims usage = subunify False+  where+    swap True x y = (y, x)+    swap False x y = (x, y)++    subunify ord bound bcs t1' t2' = do+      constraints <- getConstraints++      let nonrigid v = isNonRigid v constraints++          failure = matchError (srclocOf usage) mempty bcs t1' t2'++          unifyTypeArg bcs' (TypeArgDim d1) (TypeArgDim d2) =+            onDims' bcs' (swap ord d1 d2)+          unifyTypeArg bcs' (TypeArgType t) (TypeArgType arg_t) =+            subunify ord bound bcs' t arg_t+          unifyTypeArg _ _ _ = failure++          onDims' bcs' (d1, d2) =+            onDims+              bcs'+              bound+              nonrigid+              (applySubst (`lookupSubst` constraints) d1)+              (applySubst (`lookupSubst` constraints) d2)++      -- The types are structurally identical, as this has already+      -- been verified by the unsized type checker - we are here only+      -- to unify their sizes (and check consumption and uniqueness+      -- for functions). The 'failure' cases can be reached when the+      -- types contain distinct abstract types that the unsized+      -- checking could not distinguish, and serve as a backstop for+      -- anything it may have missed.+      case (t1', t2') of+        (Scalar (Prim pt1), Scalar (Prim pt2))+          | pt1 == pt2 -> pure ()+        ( Scalar (Record fs),+          Scalar (Record arg_fs)+          )+            | M.keys fs == M.keys arg_fs ->+                forM_ (M.toList $ M.intersectionWith (,) fs arg_fs) $ \(f, (t1, t2)) ->+                  subunify ord bound (matchingField f <> bcs) t1 t2+        ( Scalar (Sum cs),+          Scalar (Sum arg_cs)+          )+            | M.keys cs == M.keys arg_cs,+              fmap length cs == fmap length arg_cs ->+                forM_ (M.toList $ M.intersectionWith (,) cs arg_cs) $ \(c, (ts1, ts2)) ->+                  zipWithM_ (subunify ord bound (matchingConstructor c <> bcs)) ts1 ts2+        ( Scalar (TypeVar _ (QualName _ tn) targs),+          Scalar (TypeVar _ (QualName _ arg_tn) arg_targs)+          )+            | tn == arg_tn,+              length targs == length arg_targs -> do+                let bcs' = matching "When matching type arguments." <> bcs+                zipWithM_ (unifyTypeArg bcs') targs arg_targs+        ( Scalar (Arrow _ p1 d1 a1 (RetType b1_dims b1)),+          Scalar (Arrow _ p2 d2 a2 (RetType b2_dims b2))+          )+            | uncurry (<) $ swap ord d1 d2 -> do+                unifyError usage mempty bcs . withIndexLink "unify-consuming-param" $+                  "Parameters"+                    </> indent 2 (pretty d1 <> pretty a1)+                    </> "and"+                    </> indent 2 (pretty d2 <> pretty a2)+                    </> "are incompatible regarding consuming their arguments."+            | uncurry (<) $ swap ord (uniqueness b2) (uniqueness b1) -> do+                unifyError usage mempty bcs $+                  "Return types"+                    </> indent 2 (pretty b1)+                    </> "and"+                    </> indent 2 (pretty b2)+                    </> "have incompatible uniqueness."+            | otherwise -> do+                -- Introduce the existentials as size variables so they+                -- are subject to unification.  We will remove them again+                -- afterwards.+                let (r1, r2) =+                      swap+                        ord+                        (Size Nothing $ Usage Nothing mempty)+                        (UnknownSize mempty RigidUnify)+                lvl <- curLevel+                modifyConstraints (M.fromList (map (,(lvl, r1)) b1_dims) <>)+                modifyConstraints (M.fromList (map (,(lvl, r2)) b2_dims) <>)++                let bound' = bound <> mapMaybe pname [p1, p2] <> b1_dims <> b2_dims+                subunify+                  (not ord)+                  bound+                  (matching "When matching parameter types." <> bcs)+                  a1+                  a2+                subunify+                  ord+                  bound'+                  (matching "When matching return types." <> bcs)+                  (toStruct b1')+                  (toStruct b2')++                -- If a flexible existential size was resolved to a pending+                -- instantiated size, then that size is existential. This is how+                -- a hole absorbs an existential size from a type it is unified+                -- with. See Note [Size Inference] in+                -- Language.Futhark.TypeChecker.Terms.+                constraints_after <- getConstraints+                -- An existential that was already registered as a+                -- rigid unknown size before we made it unifiable+                -- above (e.g. a pending instantiated size bound by+                -- checkApply) stands for a size that is actually+                -- computed somewhere, so absorbing it incurs a+                -- causality obligation.+                let rigidPre d = case snd <$> M.lookup d constraints of+                      Just (UnknownSize dloc _)+                        | dloc == mempty -> Just $ locOf usage+                        | otherwise -> Just dloc+                      _ -> Nothing+                    existentialise d v usage' =+                      modifyConstraints $+                        M.adjust (fmap $ const $ ExistentialSize (Just d) (rigidPre d) usage') v+                    absorbExt d+                      | Just (Size (Just de) _) <- snd <$> M.lookup d constraints_after,+                        Var de_v _ _ <- applySubst (`lookupSubst` constraints_after) de =+                          case snd <$> M.lookup (qualLeaf de_v) constraints_after of+                            Just (InstSize l usage')+                              | l /= Unlifted ->+                                  existentialise d (qualLeaf de_v) usage'+                            Just (CopySize c _ usage')+                              | Just (InstSize l _) <- snd <$> M.lookup c constraints_after,+                                l /= Unlifted ->+                                  existentialise d c usage'+                            _ -> pure ()+                      | otherwise = pure ()+                mapM_ absorbExt (b1_dims <> b2_dims)++                -- Delete the size variables we introduced to represent the+                -- existential sizes.+                modifyConstraints $ \m -> L.foldl' (flip M.delete) m (b1_dims <> b2_dims)+            where+              (b1', b2') =+                -- Replace one parameter name with the other in the+                -- return type, in case of dependent types.  I.e.,+                -- we want type '(n: i32) -> [n]i32' to unify with+                -- type '(x: i32) -> [x]i32'.+                case (p1, p2) of+                  (Named p1', Named p2') ->+                    let f v+                          | v == p2' = Just $ ExpSubst $ sizeFromName (qualName p1') mempty+                          | otherwise = Nothing+                     in (b1, applySubst f b2)+                  (_, _) ->+                    (b1, b2)++              pname (Named x) = Just x+              pname Unnamed = Nothing+        ( Array _ (Shape (t1_d : t1_ds)) t1_et,+          Array _ (Shape (t2_d : t2_ds)) t2_et+          ) -> do+            onDims' bcs (swap ord t1_d t2_d)+            subunify+              ord+              bound+              bcs+              (arrayOf (Shape t1_ds) (Scalar t1_et))+              (arrayOf (Shape t2_ds) (Scalar t2_et))+        _ -> failure++anyBound :: [VName] -> ExpBase Info VName -> Bool+anyBound bound e = any (`S.member` fvVars (freeInExp e)) bound++unifySizes :: (MonadUnify m) => Usage -> UnifySizes m+unifySizes usage bcs bound nonrigid e1 e2+  | Just es <- similarExps e1 e2 =+      mapM_ (uncurry $ unifySizes usage bcs bound nonrigid) es+unifySizes usage bcs bound nonrigid (Var v1 _ _) e2+  | Just lvl1 <- nonrigid (qualLeaf v1),+    not (anyBound bound e2) || (qualLeaf v1 `elem` bound) =+      linkVarToDim usage bcs (qualLeaf v1) lvl1 e2+unifySizes usage bcs bound nonrigid e1 (Var v2 _ _)+  | Just lvl2 <- nonrigid (qualLeaf v2),+    not (anyBound bound e1) || (qualLeaf v2 `elem` bound) =+      linkVarToDim usage bcs (qualLeaf v2) lvl2 e1+unifySizes usage bcs bound _ e1 e2 = do+  -- A size arising from a type parameter instantiation may be linked+  -- to sizes bound within the instantiated type itself (reconstructing+  -- a dependent function type), and when it meets any other bound+  -- size (an existential), it is determined to be existential itself,+  -- rather than this being an error. This is the only way we can know+  -- how instantiated sizes depend on binders and existentials. See+  -- Note [Size Inference] in Language.Futhark.TypeChecker.Terms.+  linked <- (||) <$> maybeLocalLink e1 e2 <*> maybeLocalLink e2 e1+  absorbed <-+    if linked+      then pure True+      else (||) <$> maybeAbsorb e1 e2 <*> maybeAbsorb e2 e1+  unless absorbed $ do+    notes <- (<>) <$> dimNotes usage e1 <*> dimNotes usage e2+    anon1 <- instMeetsAnonymous e1 e2+    anon2 <- instMeetsAnonymous e2 e1+    if anon1 || anon2+      then+        unifyError usage notes bcs $+          "Sizes"+            <+> dquotes (pretty e1)+            <+> "and"+            <+> dquotes (pretty e2)+            <+> "do not match."+            </> textwrap "This is because a type parameter would be instantiated with a type containing anonymous sizes."+      else+        unifyError usage notes bcs $+          "Sizes"+            <+> dquotes (pretty e1)+            <+> "and"+            <+> dquotes (pretty e2)+            <+> "do not match."+  where+    instConstraint constraints v = do+      c <- snd <$> M.lookup v constraints+      case c of+        InstSize {} -> Just c+        CopySize {} -> Just c+        ExistentialSize {} -> Just c+        _ -> Nothing+    -- If the absorbed size is a rigid unknown size, then the+    -- existential stands for a size that is actually computed+    -- somewhere, and uses of it are subject to the causality check.+    -- Sizes bound in the type itself (existentials of a declared+    -- type, parameters) carry no such obligation.+    existentialise v other usage' = do+      constraints <- getConstraints+      let rigidLoc w = case snd <$> M.lookup w constraints of+            Just (UnknownSize wloc _)+              | wloc == mempty -> Just $ locOf usage+              | otherwise -> Just wloc+            _ -> Nothing+          computed_at =+            listToMaybe $ mapMaybe rigidLoc $ S.toList $ fvVars $ freeInExp other+      modifyConstraints $ M.adjust (fmap $ const $ ExistentialSize key computed_at usage') v+      where+        key = case other of+          Var other_v _ _ -> Just $ qualLeaf other_v+          _ -> Nothing+    -- Linking is fine if every bound size mentioned is a binder of+    -- the instantiated type itself (a registered 'ParamSize'), as+    -- instantiated size variables occur exactly once, and binders+    -- are cloned between occurrences of the instantiated type.+    maybeLocalLink (Var v _ _) other+      | anyBound bound other,+        qualLeaf v `notElem` bound = do+          constraints <- getConstraints+          let mentioned = filter (`elem` bound) $ S.toList $ fvVars $ freeInExp other+              registeredBinder bv = case snd <$> M.lookup bv constraints of+                Just (ParamSize _) -> True+                _ -> False+          case instConstraint constraints (qualLeaf v) of+            Just c+              | all registeredBinder mentioned,+                notExistential c -> do+                  modifyConstraints $+                    M.adjust (fmap $ const $ Size (Just other) usage) (qualLeaf v)+                  pure True+            _ -> pure False+      where+        notExistential ExistentialSize {} = False+        notExistential _ = True+    maybeLocalLink _ _ = pure False+    maybeAbsorb (Var v _ _) other+      | anyBound bound other,+        qualLeaf v `notElem` bound = do+          constraints <- getConstraints+          case snd <$> M.lookup (qualLeaf v) constraints of+            Just (InstSize l usage')+              | l /= Unlifted ->+                  True <$ existentialise (qualLeaf v) other usage'+            Just ExistentialSize {} ->+              pure True+            Just (CopySize c _ usage') ->+              case snd <$> M.lookup c constraints of+                Just (InstSize l _)+                  | l /= Unlifted -> True <$ existentialise c other usage'+                Just ExistentialSize {} -> pure True+                _ -> pure False+            _ -> pure False+    maybeAbsorb _ _ = pure False+    instMeetsAnonymous (Var v _ _) other+      | anyBound bound other = do+          constraints <- getConstraints+          pure $ isJust $ instConstraint constraints $ qualLeaf v+    instMeetsAnonymous _ _ = pure False++-- | Unifies two types.+unify :: (MonadUnify m) => Usage -> StructType -> StructType -> m ()+unify usage = unifyWith (unifySizes usage) usage mempty mempty++linkVarToDim ::+  (MonadUnify m) =>+  Usage ->+  BreadCrumbs ->+  VName ->+  Level ->+  Exp ->+  m ()+linkVarToDim usage bcs vn lvl e = do+  constraints <- getConstraints++  -- A copy of an instantiated size is equal to its canonical+  -- variable as long as the size is not existential, so links are+  -- expressed in terms of canonical variables: both when the linked+  -- variable is a copy, and when copies occur in the expression+  -- linked to.+  let canonize v = case snd <$> M.lookup v constraints of+        Just (CopySize c _ _) ->+          Just $ ExpSubst $ sizeFromName (qualName c) $ srclocOf usage+        _ -> Nothing+      e' = applySubst canonize e++  case snd <$> M.lookup vn constraints of+    Just (CopySize c _ _)+      | Just (c_lvl, _) <- M.lookup c constraints ->+          linkVarToDim usage bcs c c_lvl e'+    _+      -- Linking a size to itself is a no-op. This can occur when+      -- unifying a canonical size with one of its own copies.+      | Var (QualName _ e_v) _ _ <- e',+        e_v == vn ->+          pure ()+      | otherwise -> do+          mapM_ (checkVar constraints) $ fvVars $ freeInExp e'++          modifyConstraints $ M.insert vn (lvl, Size (Just e') usage)+  where+    checkVar _ dim'+      | vn == dim' = do+          notes <- dimNotes usage e+          unifyError usage notes bcs . withIndexLink "occurs-check" $+            "Occurs check: cannot instantiate"+              <+> dquotes (prettyName vn)+              <+> "with"+              <+> dquotes (pretty e)+              <+> "."+    checkVar constraints dim'+      | Just (dim_lvl, c) <- dim' `M.lookup` constraints,+        dim_lvl >= lvl =+          case c of+            ParamSize {} -> do+              notes <- dimNotes usage e+              unifyError usage notes bcs $+                withIndexLink "scope-violation" $+                  "Cannot link size"+                    <+> dquotes (prettyName vn)+                    <+> "to"+                    <+> dquotes (pretty e)+                    <+> "(scope violation)."+                    </> "This is because"+                    <+> dquotes (pretty $ qualName dim')+                    <+> "is not in scope when"+                    <+> dquotes (prettyName vn)+                    <+> "is introduced."+            _ -> modifyConstraints $ M.insert dim' (lvl, c)+    checkVar _ _ = pure ()++newDimOnMismatch ::+  (MonadUnify m) =>+  Loc ->+  StructType ->+  StructType ->+  m (StructType, [VName])+newDimOnMismatch loc t1 t2 = do+  (t, seen) <- runStateT (matchDims onDims t1 t2) mempty+  pure (t, M.elems seen)+  where+    r = RigidCond t1 t2+    same (e1, e2) =+      maybe False (all same) $ similarExps e1 e2+    onDims _ d1 d2+      | same (d1, d2) = pure d1+      | otherwise = do+          -- Remember mismatches we have seen before and reuse the+          -- same new size.+          maybe_d <- gets $ M.lookup (d1, d2)+          case maybe_d of+            Just d -> pure $ sizeFromName (qualName d) $ srclocOf loc+            Nothing -> do+              d <- lift $ newRigidDim loc r "differ"+              modify $ M.insert (d1, d2) d+              pure $ sizeFromName (qualName d) $ srclocOf loc++-- | Like unification, but creates new size variables where mismatches+-- occur.  Returns the new dimensions thus created.+unifyMostCommon ::+  (MonadUnify m) =>+  Usage ->+  StructType ->+  StructType ->+  m (StructType, [VName])+unifyMostCommon usage t1 t2 = do+  -- Like 'unifySizes', except we do not fail on mismatches - these+  -- are instead turned into fresh existential sizes in+  -- 'newDimOnMismatch'. The most annoying thing is that we have to+  -- replicate scope checking, because we don't want to link if it+  -- would fail.+  constraints <- getConstraints++  let expFreeVars = fvVars . freeInExp+      varLevel v = fst <$> M.lookup v constraints++      -- Check that linking to this expression would not fail in linkVarToDim+      -- due to a ParamSize at a level >= the target level. This replicates the+      -- scope check performed by linkVarToDim's checkVar.+      wouldFail lvl v =+        case M.lookup v constraints of+          Just (dim_lvl, ParamSize {}) -> dim_lvl >= lvl+          _ -> False++      -- Can we link a variable at the given level to an expression with the+      -- given free variables? FIXME: something her is fishy. Why do we need to+      -- treat ParamSize specially in wouldFail? Why is the level check for the+      -- other variables not enough?+      canLink lvl vn bound fvs =+        L.foldl' max 0 (mapMaybe varLevel $ S.toList fvs) <= lvl+          && not (any (`S.member` fvs) bound)+          && not (any (wouldFail lvl) $ S.toList fvs)+          && not (vn `S.member` fvs)++      onDims bcs bound nonrigid e1 e2+        | Just es <- similarExps e1 e2 =+            mapM_ (uncurry $ onDims bcs bound nonrigid) es+      onDims bcs bound nonrigid (Var v1 _ _) e2+        | Just lvl1 <- nonrigid (qualLeaf v1),+          canLink lvl1 (qualLeaf v1) bound (expFreeVars e2) =+            linkVarToDim usage bcs (qualLeaf v1) lvl1 e2+      onDims bcs bound nonrigid e1 (Var v2 _ _)+        | Just lvl2 <- nonrigid (qualLeaf v2),+          canLink lvl2 (qualLeaf v2) bound (expFreeVars e1) =+            linkVarToDim usage bcs (qualLeaf v2) lvl2 e1+      onDims _ _ _ _ _ = pure ()++  unifyWith onDims usage mempty mempty t1 t2+  t1' <- normTypeFully t1+  t2' <- normTypeFully t2+  newDimOnMismatch (locOf usage) t1' t2'++-- | Replace *all* dimensions with distinct fresh size variables.+allDimsFreshInType ::+  (MonadUnify m) =>+  Usage ->+  Rigidity ->+  Name ->+  TypeBase d als ->+  m (TypeBase Size als, M.Map VName d)+allDimsFreshInType usage r desc t =+  runStateT (bitraverse onDim pure t) mempty+  where+    onDim d = do+      v <- lift $ newDimVar usage r desc+      modify $ M.insert v d+      pure $ sizeFromName (qualName v) $ srclocOf usage++-- Simple pure MonadUnify implementation for unification outside of+-- the term checker. The constraints contain only sizes.++type UnifyMState = (Constraints, Int)++newtype UnifyM a = UnifyM (StateT UnifyMState (Except TypeError) a)+  deriving+    ( Monad,+      Functor,+      Applicative,+      MonadState UnifyMState,+      MonadError TypeError+    )++instance MonadUnify UnifyM where+  getConstraints = gets fst+  putConstraints x = modify $ \(_, i) -> (x, i)++  newDimVar usage rigidity name = do+    (x, i) <- get+    put (x, i + 1)+    -- Note that the level is 1, so that fresh sizes may be linked to+    -- the rigid parameters, which are at level 0.+    let dim = VName (mkTypeVarName name i) i+    case rigidity of+      Rigid src ->+        modifyConstraints $+          M.insert dim (1, UnknownSize (locOf usage) src)+      Nonrigid ->+        modifyConstraints $+          M.insert dim (1, Size Nothing usage)+    pure dim++  curLevel = pure 1++  unifyError loc notes bcs doc =+    throwError $ TypeError (locOf loc) notes $ doc <> pretty bcs++  matchError loc notes bcs t1 t2 =+    throwError $ TypeError (locOf loc) notes $ doc <> pretty bcs+    where+      doc =+        "Types"+          </> indent 2 (pretty t1)+          </> "and"+          </> indent 2 (pretty t2)+          </> "do not match."++runUnifyM :: [TypeParam] -> [TypeParam] -> UnifyM a -> Either TypeError a+runUnifyM rigid_tparams nonrigid_tparams (UnifyM m) =+  runExcept $ evalStateT m (constraints, 0)+  where+    constraints =+      M.fromList $+        mapMaybe nonrigid nonrigid_tparams <> mapMaybe rigid rigid_tparams+    nonrigid (TypeParamDim p ploc) =+      Just (p, (1, Size Nothing $ Usage Nothing $ locOf ploc))+    nonrigid TypeParamType {} = Nothing+    rigid (TypeParamDim p ploc) = Just (p, (0, ParamSize $ locOf ploc))+    rigid TypeParamType {} = Nothing++-- | Check that two types match, instantiating the nonrigid type+-- parameters of the second type as necessary. This is used when+-- matching a value or type in a module against a specification.+--+-- This works in two phases. First the types are checked while+-- disregarding sizes entirely, using the same constraint solver as+-- the unsized type checker. This also determines the instantiation+-- of the nonrigid type parameters, up to sizes. Then the+-- instantiations, given fresh size variables, are substituted into+-- the second type, and the sizes are checked with ordinary+-- (size-only) unification.+doUnification ::+  Loc ->+  [TypeParam] ->+  [TypeParam] ->+  StructType ->+  StructType ->+  Either TypeError ()+doUnification loc rigid_tparams nonrigid_tparams spec_t t = do+  -- Phase 1: types.+  let typarams =+        M.fromList+          [ (v, (0, l, locOf tploc))+          | TypeParamType l v tploc <- rigid_tparams+          ]+      tyvars =+        M.fromList+          [ (v, (1, TyVarFree (locOf tploc) l))+          | TypeParamType l v tploc <- nonrigid_tparams+          ]+      ct = CtEq (Reason loc) (unsized spec_t) (unsized t)+  (_, solution) <- TySolve.solve [ct] typarams tyvars++  -- The solver does not verify that instantiations respect the+  -- liftedness of the instantiated type parameter, so we check that+  -- here.+  mapM_ (checkLiftedness solution) nonrigid_tparams++  -- Phase 2: sizes.+  runUnifyM rigid_tparams nonrigid_tparams $ do+    -- Give the instantiations of the type parameters fresh size+    -- variables. Crucially, each type parameter is instantiated only+    -- once, so multiple occurrences of the same type parameter will+    -- have the same sizes.+    substs <- fmap (M.fromList . catMaybes) . forM nonrigid_tparams $ \case+      TypeParamType _ v _+        | Just (Right sol_t) <- M.lookup v solution -> do+            (sol_t', _) <-+              allDimsFreshInType (Usage Nothing loc) Nonrigid "d" sol_t+            pure $ Just (v, Subst [] $ RetType [] sol_t')+      _ -> pure Nothing+    unify (Usage Nothing loc) spec_t $ applySubst (`M.lookup` substs) t+  where+    unsized = first $ const ()++    rigid_liftedness =+      M.fromList [(v, l) | TypeParamType l v _ <- rigid_tparams]++    checkLiftedness _ (TypeParamDim {}) = pure ()+    checkLiftedness _ (TypeParamType Lifted _ _) = pure ()+    checkLiftedness solution (TypeParamType l v _)+      | Just (Right inst_t) <- M.lookup v solution = do+          unless (orderZero inst_t) . Left . TypeError loc mempty $+            "Cannot instantiate type parameter"+              <+> dquotes (prettyName v)+              <+> "with functional type"+              </> indent 2 (pretty inst_t)+          case mapMaybe badParam $ S.toList $ typeVars inst_t of+            v' : _ ->+              Left . TypeError loc mempty $+                "Cannot instantiate type parameter"+                  <+> dquotes (prettyName v)+                  <+> "with type containing lifted type parameter"+                  <+> dquotes (prettyName v')+                  <> "."+            [] -> pure ()+      | otherwise = pure ()+      where+        badParam v' = do+          l' <- M.lookup v' rigid_liftedness+          guard $ case l of+            Unlifted -> l' /= Unlifted+            _ -> l' == Lifted+          Just v'
+ src/Language/Futhark/TypeChecker/UnionFind.hs view
@@ -0,0 +1,137 @@+module Language.Futhark.TypeChecker.UnionFind+  ( TyVarNode,+    TyVarSol (..),+    makeTyVarNode,+    makeTyParamNode,+    find,+    getSol,+    getKey,+    assignNewSol,+    union,+    unionNewSol,+  )+where++import Control.Monad (when)+import Control.Monad.ST (ST)+import Data.STRef+  ( STRef,+    modifySTRef',+    newSTRef,+    readSTRef,+    writeSTRef,+  )+import Language.Futhark (Liftedness, Loc)+import Language.Futhark.TypeChecker.Constraints+  ( CtType,+    Level,+    TyVar,+    TyVarInfo,+  )++type Type = CtType ()++-- | A (partial) solution for a type variable.+data TyVarSol+  = -- | Has been assigned this type.+    Solved Type+  | -- | Is an explicit (rigid) type parameter in the source program.+    Param Level Liftedness Loc+  | -- | Is unsolved but has this constraint.+    Unsolved (TyVarInfo ())+  deriving (Show, Eq)++-- | A node in the union-find graph containing information about a type+-- variable.+newtype TyVarNode s = Node (STRef s (NodeInfo s)) deriving (Eq)++data NodeInfo s+  = Link !(TyVarNode s)+  | Repr !ReprInfo++data ReprInfo = ReprInfo+  { solution :: !TyVarSol,+    key :: !TyVar+  }++-- | Create a fresh node of a type variable and return it. A fresh node+-- is in the equivalence class that contains only itself.+makeTyVarNode :: TyVar -> TyVarInfo () -> ST s (TyVarNode s)+makeTyVarNode tv constraint = do+  let r =+        ReprInfo+          { solution = Unsolved constraint,+            key = tv+          }+  ref <- newSTRef $ Repr r+  pure $ Node ref++-- | Create a fresh node of a type parameter and return it. A fresh node+-- is in the equivalence class that contains only itself.+makeTyParamNode :: TyVar -> Level -> Liftedness -> Loc -> ST s (TyVarNode s)+makeTyParamNode tv lvl lft loc = do+  let r =+        ReprInfo+          { solution = Param lvl lft loc,+            key = tv+          }+  ref <- newSTRef $ Repr r+  pure $ Node ref++-- | @find node@ returns the representative of @node@'s+-- equivalence class and the information associated with+-- this equivalence class.+--+-- This method performs the path compression.+find :: TyVarNode s -> ST s (TyVarNode s, ReprInfo)+find node@(Node ref) = do+  node_info <- readSTRef ref+  case node_info of+    -- Input node is representative.+    Repr repr_info -> pure (node, repr_info)+    -- Input node's parent is another node.+    Link parent -> do+      a@(repr, _) <- find parent+      when (repr /= parent) $+        -- Performing path compression.+        writeSTRef ref $+          Link repr+      pure a++-- | Return the solution associated with the argument node's+-- equivalence class.+getSol :: TyVarNode s -> ST s TyVarSol+getSol node = solution . snd <$> find node++-- | Return the name of the representative type variable.+getKey :: TyVarNode s -> ST s TyVar+getKey node = key . snd <$> find node++-- | Assign a new solution/type to the node's equivalence class.+--+-- Precondition: The node is in an equivalence class representing an+-- unsolved/flexible type variable.+assignNewSol :: TyVarNode s -> TyVarSol -> ST s ()+assignNewSol node new_sol = do+  (Node ref, repr_info) <- find node+  modifySTRef' ref $ const . Repr $ repr_info {solution = new_sol}++-- | Join the equivalence classes of the nodes. The resulting equivalence+-- class has the same solution and key as the second argument.+union :: TyVarNode s -> TyVarNode s -> ST s ()+union n1 n2 = do+  Node ref <- fst <$> find n1+  root2 <- fst <$> find n2++  writeSTRef ref $ Link root2++-- | Join the equivalence classes of the nodes. The resulting equivalence+-- class has the same key as the second argument while @new_sol@ is the+-- new solution.+unionNewSol :: TyVarNode s -> TyVarNode s -> TyVarSol -> ST s ()+unionNewSol n1 n2 new_sol = do+  Node ref1 <- fst <$> find n1+  (root2@(Node ref2), repr_info) <- find n2++  modifySTRef' ref2 $ const . Repr $ repr_info {solution = new_sol}+  writeSTRef ref1 $ Link root2