diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/docs/conf.py b/docs/conf.py --- a/docs/conf.py +++ b/docs/conf.py @@ -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", diff --git a/docs/error-index.rst b/docs/error-index.rst --- a/docs/error-index.rst +++ b/docs/error-index.rst @@ -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 `, 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 ------------ diff --git a/docs/glossary.rst b/docs/glossary.rst --- a/docs/glossary.rst +++ b/docs/glossary.rst @@ -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` *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`. - Parametric module A function from :term:`modules` 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`. + 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 `. + 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 `. 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 `. When something is nested inside multiple + loops, it may be variant to just one of them. diff --git a/docs/index.rst b/docs/index.rst --- a/docs/index.rst +++ b/docs/index.rst @@ -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 diff --git a/docs/language-reference.rst b/docs/language-reference.rst --- a/docs/language-reference.rst +++ b/docs/language-reference.rst @@ -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 diff --git a/docs/man/futhark-bench.rst b/docs/man/futhark-bench.rst --- a/docs/man/futhark-bench.rst +++ b/docs/man/futhark-bench.rst @@ -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)` diff --git a/docs/man/futhark-benchcmp.rst b/docs/man/futhark-benchcmp.rst new file mode 100644 --- /dev/null +++ b/docs/man/futhark-benchcmp.rst @@ -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)` diff --git a/docs/man/futhark.rst b/docs/man/futhark.rst --- a/docs/man/futhark.rst +++ b/docs/man/futhark.rst @@ -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)` diff --git a/docs/versus-other-languages.rst b/docs/versus-other-languages.rst --- a/docs/versus-other-languages.rst +++ b/docs/versus-other-languages.rst @@ -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 ----- diff --git a/futhark.cabal b/futhark.cabal --- a/futhark.cabal +++ b/futhark.cabal @@ -1,6 +1,6 @@ cabal-version: 3.0 name: futhark -version: 0.26.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 diff --git a/prelude/soacs.fut b/prelude/soacs.fut --- a/prelude/soacs.fut +++ b/prelude/soacs.fut @@ -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.. [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) diff --git a/rts/c/atomics32.h b/rts/c/atomics32.h --- a/rts/c/atomics32.h +++ b/rts/c/atomics32.h @@ -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 diff --git a/rts/c/atomics64.h b/rts/c/atomics64.h --- a/rts/c/atomics64.h +++ b/rts/c/atomics64.h @@ -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 diff --git a/rts/c/backends/opencl.h b/rts/c/backends/opencl.h --- a/rts/c/backends/opencl.h +++ b/rts/c/backends/opencl.h @@ -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 ", diff --git a/rts/c/scheduler.h b/rts/c/scheduler.h --- a/rts/c/scheduler.h +++ b/rts/c/scheduler.h @@ -101,6 +101,10 @@ #include #include #include +#elif defined(__OpenBSD__) +#include +#include +#include #elif defined(__EMSCRIPTEN__) #include #include @@ -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 diff --git a/rts/opencl/prelude.cl b/rts/opencl/prelude.cl --- a/rts/opencl/prelude.cl +++ b/rts/opencl/prelude.cl @@ -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;"); } diff --git a/rts/python/opencl.py b/rts/python/opencl.py --- a/rts/python/opencl.py +++ b/rts/python/opencl.py @@ -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 += [ diff --git a/src-testing/Futhark/IR/Syntax/CoreTests.hs b/src-testing/Futhark/IR/Syntax/CoreTests.hs --- a/src-testing/Futhark/IR/Syntax/CoreTests.hs +++ b/src-testing/Futhark/IR/Syntax/CoreTests.hs @@ -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 ] diff --git a/src-testing/Generated/AllFutBenchmarks.hs b/src-testing/Generated/AllFutBenchmarks.hs new file mode 100644 --- /dev/null +++ b/src-testing/Generated/AllFutBenchmarks.hs @@ -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) + ] diff --git a/src-testing/Generated/AllFutBenchmarks/Accelerate/Nbody/Nbodybh.hs b/src-testing/Generated/AllFutBenchmarks/Accelerate/Nbody/Nbodybh.hs new file mode 100644 --- /dev/null +++ b/src-testing/Generated/AllFutBenchmarks/Accelerate/Nbody/Nbodybh.hs @@ -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))] + ) + ] diff --git a/src-testing/Language/Futhark/SyntaxTests.hs b/src-testing/Language/Futhark/SyntaxTests.hs --- a/src-testing/Language/Futhark/SyntaxTests.hs +++ b/src-testing/Language/Futhark/SyntaxTests.hs @@ -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 = diff --git a/src-testing/Language/Futhark/TypeChecker/TySolveBenchmarks.hs b/src-testing/Language/Futhark/TypeChecker/TySolveBenchmarks.hs new file mode 100644 --- /dev/null +++ b/src-testing/Language/Futhark/TypeChecker/TySolveBenchmarks.hs @@ -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] diff --git a/src-testing/Language/Futhark/TypeChecker/TySolveTests.hs b/src-testing/Language/Futhark/TypeChecker/TySolveTests.hs new file mode 100644 --- /dev/null +++ b/src-testing/Language/Futhark/TypeChecker/TySolveTests.hs @@ -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")]) + ] diff --git a/src-testing/Language/Futhark/TypeChecker/TypesTests.hs b/src-testing/Language/Futhark/TypeChecker/TypesTests.hs --- a/src-testing/Language/Futhark/TypeChecker/TypesTests.hs +++ b/src-testing/Language/Futhark/TypeChecker/TypesTests.hs @@ -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 = diff --git a/src-testing/Language/Futhark/TypeCheckerTests.hs b/src-testing/Language/Futhark/TypeCheckerTests.hs --- a/src-testing/Language/Futhark/TypeCheckerTests.hs +++ b/src-testing/Language/Futhark/TypeCheckerTests.hs @@ -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 ] diff --git a/src-testing/futhark_benchmarks.hs b/src-testing/futhark_benchmarks.hs --- a/src-testing/futhark_benchmarks.hs +++ b/src-testing/futhark_benchmarks.hs @@ -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 ] diff --git a/src/Futhark.hs b/src/Futhark.hs --- a/src/Futhark.hs +++ b/src/Futhark.hs @@ -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 diff --git a/src/Futhark/AD/Fwd.hs b/src/Futhark/AD/Fwd.hs --- a/src/Futhark/AD/Fwd.hs +++ b/src/Futhark/AD/Fwd.hs @@ -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 diff --git a/src/Futhark/AD/Rev.hs b/src/Futhark/AD/Rev.hs --- a/src/Futhark/AD/Rev.hs +++ b/src/Futhark/AD/Rev.hs @@ -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 diff --git a/src/Futhark/AD/Rev/Acc.hs b/src/Futhark/AD/Rev/Acc.hs --- a/src/Futhark/AD/Rev/Acc.hs +++ b/src/Futhark/AD/Rev/Acc.hs @@ -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 diff --git a/src/Futhark/AD/Rev/Map.hs b/src/Futhark/AD/Rev/Map.hs --- a/src/Futhark/AD/Rev/Map.hs +++ b/src/Futhark/AD/Rev/Map.hs @@ -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. diff --git a/src/Futhark/AD/Rev/Monad.hs b/src/Futhark/AD/Rev/Monad.hs --- a/src/Futhark/AD/Rev/Monad.hs +++ b/src/Futhark/AD/Rev/Monad.hs @@ -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 diff --git a/src/Futhark/AD/Rev/SOAC.hs b/src/Futhark/AD/Rev/SOAC.hs --- a/src/Futhark/AD/Rev/SOAC.hs +++ b/src/Futhark/AD/Rev/SOAC.hs @@ -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 diff --git a/src/Futhark/AD/Shared.hs b/src/Futhark/AD/Shared.hs --- a/src/Futhark/AD/Shared.hs +++ b/src/Futhark/AD/Shared.hs @@ -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) => diff --git a/src/Futhark/Analysis/Alias.hs b/src/Futhark/Analysis/Alias.hs --- a/src/Futhark/Analysis/Alias.hs +++ b/src/Futhark/Analysis/Alias.hs @@ -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 diff --git a/src/Futhark/Analysis/CallGraph.hs b/src/Futhark/Analysis/CallGraph.hs --- a/src/Futhark/Analysis/CallGraph.hs +++ b/src/Futhark/Analysis/CallGraph.hs @@ -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 = diff --git a/src/Futhark/Analysis/DataDependencies.hs b/src/Futhark/Analysis/DataDependencies.hs --- a/src/Futhark/Analysis/DataDependencies.hs +++ b/src/Futhark/Analysis/DataDependencies.hs @@ -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 = diff --git a/src/Futhark/Analysis/Metrics.hs b/src/Futhark/Analysis/Metrics.hs --- a/src/Futhark/Analysis/Metrics.hs +++ b/src/Futhark/Analysis/Metrics.hs @@ -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 diff --git a/src/Futhark/CLI/Benchcmp.hs b/src/Futhark/CLI/Benchcmp.hs --- a/src/Futhark/CLI/Benchcmp.hs +++ b/src/Futhark/CLI/Benchcmp.hs @@ -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 () [] " " f +main = mainWithOptions defaultConfig options " " 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 diff --git a/src/Futhark/CLI/Dev.hs b/src/Futhark/CLI/Dev.hs --- a/src/Futhark/CLI/Dev.hs +++ b/src/Futhark/CLI/Dev.hs @@ -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 [], diff --git a/src/Futhark/CodeGen/Backends/GenericPython.hs b/src/Futhark/CodeGen/Backends/GenericPython.hs --- a/src/Futhark/CodeGen/Backends/GenericPython.hs +++ b/src/Futhark/CodeGen/Backends/GenericPython.hs @@ -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 diff --git a/src/Futhark/CodeGen/Backends/GenericPython/AST.hs b/src/Futhark/CodeGen/Backends/GenericPython/AST.hs --- a/src/Futhark/CodeGen/Backends/GenericPython/AST.hs +++ b/src/Futhark/CodeGen/Backends/GenericPython/AST.hs @@ -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) diff --git a/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs b/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs --- a/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU/SegRed.hs @@ -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 diff --git a/src/Futhark/Construct.hs b/src/Futhark/Construct.hs --- a/src/Futhark/Construct.hs +++ b/src/Futhark/Construct.hs @@ -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. diff --git a/src/Futhark/IR/GPU.hs b/src/Futhark/IR/GPU.hs --- a/src/Futhark/IR/GPU.hs +++ b/src/Futhark/IR/GPU.hs @@ -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 diff --git a/src/Futhark/IR/GPU/Simplify.hs b/src/Futhark/IR/GPU/Simplify.hs --- a/src/Futhark/IR/GPU/Simplify.hs +++ b/src/Futhark/IR/GPU/Simplify.hs @@ -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 diff --git a/src/Futhark/IR/Parse.hs b/src/Futhark/IR/Parse.hs --- a/src/Futhark/IR/Parse.hs +++ b/src/Futhark/IR/Parse.hs @@ -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 $ diff --git a/src/Futhark/IR/Pretty.hs b/src/Futhark/IR/Pretty.hs --- a/src/Futhark/IR/Pretty.hs +++ b/src/Futhark/IR/Pretty.hs @@ -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) = "\\" diff --git a/src/Futhark/IR/Prop.hs b/src/Futhark/IR/Prop.hs --- a/src/Futhark/IR/Prop.hs +++ b/src/Futhark/IR/Prop.hs @@ -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 diff --git a/src/Futhark/IR/Prop/Names.hs b/src/Futhark/IR/Prop/Names.hs --- a/src/Futhark/IR/Prop/Names.hs +++ b/src/Futhark/IR/Prop/Names.hs @@ -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) diff --git a/src/Futhark/IR/Prop/Types.hs b/src/Futhark/IR/Prop/Types.hs --- a/src/Futhark/IR/Prop/Types.hs +++ b/src/Futhark/IR/Prop/Types.hs @@ -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 diff --git a/src/Futhark/IR/Rephrase.hs b/src/Futhark/IR/Rephrase.hs --- a/src/Futhark/IR/Rephrase.hs +++ b/src/Futhark/IR/Rephrase.hs @@ -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 diff --git a/src/Futhark/IR/RetType.hs b/src/Futhark/IR/RetType.hs --- a/src/Futhark/IR/RetType.hs +++ b/src/Futhark/IR/RetType.hs @@ -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 diff --git a/src/Futhark/IR/SOACS.hs b/src/Futhark/IR/SOACS.hs --- a/src/Futhark/IR/SOACS.hs +++ b/src/Futhark/IR/SOACS.hs @@ -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 diff --git a/src/Futhark/IR/SOACS/SOAC.hs b/src/Futhark/IR/SOACS/SOAC.hs --- a/src/Futhark/IR/SOACS/SOAC.hs +++ b/src/Futhark/IR/SOACS/SOAC.hs @@ -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 diff --git a/src/Futhark/IR/SOACS/Simplify.hs b/src/Futhark/IR/SOACS/Simplify.hs --- a/src/Futhark/IR/SOACS/Simplify.hs +++ b/src/Futhark/IR/SOACS/Simplify.hs @@ -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) diff --git a/src/Futhark/IR/SegOp.hs b/src/Futhark/IR/SegOp.hs --- a/src/Futhark/IR/SegOp.hs +++ b/src/Futhark/IR/SegOp.hs @@ -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 diff --git a/src/Futhark/IR/Syntax.hs b/src/Futhark/IR/Syntax.hs --- a/src/Futhark/IR/Syntax.hs +++ b/src/Futhark/IR/Syntax.hs @@ -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) diff --git a/src/Futhark/IR/Syntax/Core.hs b/src/Futhark/IR/Syntax/Core.hs --- a/src/Futhark/IR/Syntax/Core.hs +++ b/src/Futhark/IR/Syntax/Core.hs @@ -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 diff --git a/src/Futhark/IR/Traversals.hs b/src/Futhark/IR/Traversals.hs --- a/src/Futhark/IR/Traversals.hs +++ b/src/Futhark/IR/Traversals.hs @@ -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) diff --git a/src/Futhark/IR/TypeCheck.hs b/src/Futhark/IR/TypeCheck.hs --- a/src/Futhark/IR/TypeCheck.hs +++ b/src/Futhark/IR/TypeCheck.hs @@ -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 "") (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 "" 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 () diff --git a/src/Futhark/Internalise/Bindings.hs b/src/Futhark/Internalise/Bindings.hs --- a/src/Futhark/Internalise/Bindings.hs +++ b/src/Futhark/Internalise/Bindings.hs @@ -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] -> diff --git a/src/Futhark/Internalise/Defunctionalise.hs b/src/Futhark/Internalise/Defunctionalise.hs --- a/src/Futhark/Internalise/Defunctionalise.hs +++ b/src/Futhark/Internalise/Defunctionalise.hs @@ -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. diff --git a/src/Futhark/Internalise/Exps.hs b/src/Futhark/Internalise/Exps.hs --- a/src/Futhark/Internalise/Exps.hs +++ b/src/Futhark/Internalise/Exps.hs @@ -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 diff --git a/src/Futhark/Internalise/Monad.hs b/src/Futhark/Internalise/Monad.hs --- a/src/Futhark/Internalise/Monad.hs +++ b/src/Futhark/Internalise/Monad.hs @@ -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 diff --git a/src/Futhark/Internalise/Monomorphise.hs b/src/Futhark/Internalise/Monomorphise.hs --- a/src/Futhark/Internalise/Monomorphise.hs +++ b/src/Futhark/Internalise/Monomorphise.hs @@ -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. diff --git a/src/Futhark/Internalise/TypesValues.hs b/src/Futhark/Internalise/TypesValues.hs --- a/src/Futhark/Internalise/TypesValues.hs +++ b/src/Futhark/Internalise/TypesValues.hs @@ -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 -> diff --git a/src/Futhark/Optimise/ArrayShortCircuiting.hs b/src/Futhark/Optimise/ArrayShortCircuiting.hs --- a/src/Futhark/Optimise/ArrayShortCircuiting.hs +++ b/src/Futhark/Optimise/ArrayShortCircuiting.hs @@ -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) diff --git a/src/Futhark/Optimise/BlkRegTiling.hs b/src/Futhark/Optimise/BlkRegTiling.hs --- a/src/Futhark/Optimise/BlkRegTiling.hs +++ b/src/Futhark/Optimise/BlkRegTiling.hs @@ -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' diff --git a/src/Futhark/Optimise/CSE.hs b/src/Futhark/Optimise/CSE.hs --- a/src/Futhark/Optimise/CSE.hs +++ b/src/Futhark/Optimise/CSE.hs @@ -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) diff --git a/src/Futhark/Optimise/DoubleBuffer.hs b/src/Futhark/Optimise/DoubleBuffer.hs --- a/src/Futhark/Optimise/DoubleBuffer.hs +++ b/src/Futhark/Optimise/DoubleBuffer.hs @@ -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 :: diff --git a/src/Futhark/Optimise/Fusion/GraphRep.hs b/src/Futhark/Optimise/Fusion/GraphRep.hs --- a/src/Futhark/Optimise/Fusion/GraphRep.hs +++ b/src/Futhark/Optimise/Fusion/GraphRep.hs @@ -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 diff --git a/src/Futhark/Optimise/Fusion/Screma.hs b/src/Futhark/Optimise/Fusion/Screma.hs --- a/src/Futhark/Optimise/Fusion/Screma.hs +++ b/src/Futhark/Optimise/Fusion/Screma.hs @@ -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) $ diff --git a/src/Futhark/Optimise/Fusion/TryFusion.hs b/src/Futhark/Optimise/Fusion/TryFusion.hs --- a/src/Futhark/Optimise/Fusion/TryFusion.hs +++ b/src/Futhark/Optimise/Fusion/TryFusion.hs @@ -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 diff --git a/src/Futhark/Optimise/InliningDeadFun.hs b/src/Futhark/Optimise/InliningDeadFun.hs --- a/src/Futhark/Optimise/InliningDeadFun.hs +++ b/src/Futhark/Optimise/InliningDeadFun.hs @@ -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", diff --git a/src/Futhark/Optimise/Simplify/Engine.hs b/src/Futhark/Optimise/Simplify/Engine.hs --- a/src/Futhark/Optimise/Simplify/Engine.hs +++ b/src/Futhark/Optimise/Simplify/Engine.hs @@ -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 diff --git a/src/Futhark/Optimise/Simplify/Rep.hs b/src/Futhark/Optimise/Simplify/Rep.hs --- a/src/Futhark/Optimise/Simplify/Rep.hs +++ b/src/Futhark/Optimise/Simplify/Rep.hs @@ -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. diff --git a/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs b/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs --- a/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs +++ b/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs @@ -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 diff --git a/src/Futhark/Optimise/Simplify/Rules/Index.hs b/src/Futhark/Optimise/Simplify/Rules/Index.hs --- a/src/Futhark/Optimise/Simplify/Rules/Index.hs +++ b/src/Futhark/Optimise/Simplify/Rules/Index.hs @@ -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. diff --git a/src/Futhark/Optimise/Simplify/Rules/Match.hs b/src/Futhark/Optimise/Simplify/Rules/Match.hs --- a/src/Futhark/Optimise/Simplify/Rules/Match.hs +++ b/src/Futhark/Optimise/Simplify/Rules/Match.hs @@ -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 [(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] diff --git a/src/Futhark/Optimise/Sink.hs b/src/Futhark/Optimise/Sink.hs --- a/src/Futhark/Optimise/Sink.hs +++ b/src/Futhark/Optimise/Sink.hs @@ -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 (+) diff --git a/src/Futhark/Pass/ExpandAllocations.hs b/src/Futhark/Pass/ExpandAllocations.hs --- a/src/Futhark/Pass/ExpandAllocations.hs +++ b/src/Futhark/Pass/ExpandAllocations.hs @@ -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. diff --git a/src/Futhark/Pass/ExplicitAllocations.hs b/src/Futhark/Pass/ExplicitAllocations.hs --- a/src/Futhark/Pass/ExplicitAllocations.hs +++ b/src/Futhark/Pass/ExplicitAllocations.hs @@ -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 diff --git a/src/Futhark/Pass/ExtractKernels.hs b/src/Futhark/Pass/ExtractKernels.hs deleted file mode 100644 --- a/src/Futhark/Pass/ExtractKernels.hs +++ /dev/null @@ -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' diff --git a/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs b/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs deleted file mode 100644 --- a/src/Futhark/Pass/ExtractKernels/BlockedKernel.hs +++ /dev/null @@ -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)) diff --git a/src/Futhark/Pass/ExtractKernels/DistributeNests.hs b/src/Futhark/Pass/ExtractKernels/DistributeNests.hs deleted file mode 100644 --- a/src/Futhark/Pass/ExtractKernels/DistributeNests.hs +++ /dev/null @@ -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 diff --git a/src/Futhark/Pass/ExtractKernels/Distribution.hs b/src/Futhark/Pass/ExtractKernels/Distribution.hs deleted file mode 100644 --- a/src/Futhark/Pass/ExtractKernels/Distribution.hs +++ /dev/null @@ -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) diff --git a/src/Futhark/Pass/ExtractKernels/ISRWIM.hs b/src/Futhark/Pass/ExtractKernels/ISRWIM.hs deleted file mode 100644 --- a/src/Futhark/Pass/ExtractKernels/ISRWIM.hs +++ /dev/null @@ -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 diff --git a/src/Futhark/Pass/ExtractKernels/Interchange.hs b/src/Futhark/Pass/ExtractKernels/Interchange.hs deleted file mode 100644 --- a/src/Futhark/Pass/ExtractKernels/Interchange.hs +++ /dev/null @@ -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) diff --git a/src/Futhark/Pass/ExtractKernels/Intrablock.hs b/src/Futhark/Pass/ExtractKernels/Intrablock.hs deleted file mode 100644 --- a/src/Futhark/Pass/ExtractKernels/Intrablock.hs +++ /dev/null @@ -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 diff --git a/src/Futhark/Pass/ExtractKernels/StreamKernel.hs b/src/Futhark/Pass/ExtractKernels/StreamKernel.hs deleted file mode 100644 --- a/src/Futhark/Pass/ExtractKernels/StreamKernel.hs +++ /dev/null @@ -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) diff --git a/src/Futhark/Pass/ExtractKernels/ToGPU.hs b/src/Futhark/Pass/ExtractKernels/ToGPU.hs deleted file mode 100644 --- a/src/Futhark/Pass/ExtractKernels/ToGPU.hs +++ /dev/null @@ -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 diff --git a/src/Futhark/Pass/ExtractMulticore.hs b/src/Futhark/Pass/ExtractMulticore.hs --- a/src/Futhark/Pass/ExtractMulticore.hs +++ b/src/Futhark/Pass/ExtractMulticore.hs @@ -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 diff --git a/src/Futhark/Pass/Flatten.hs b/src/Futhark/Pass/Flatten.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/Flatten.hs @@ -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 #-} diff --git a/src/Futhark/Pass/Flatten/BasicOp.hs b/src/Futhark/Pass/Flatten/BasicOp.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/Flatten/BasicOp.hs @@ -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 diff --git a/src/Futhark/Pass/Flatten/Builtins.hs b/src/Futhark/Pass/Flatten/Builtins.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/Flatten/Builtins.hs @@ -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) diff --git a/src/Futhark/Pass/Flatten/Distribute.hs b/src/Futhark/Pass/Flatten/Distribute.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/Flatten/Distribute.hs @@ -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. diff --git a/src/Futhark/Pass/Flatten/General.hs b/src/Futhark/Pass/Flatten/General.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/Flatten/General.hs @@ -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 diff --git a/src/Futhark/Pass/Flatten/Incremental.hs b/src/Futhark/Pass/Flatten/Incremental.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/Flatten/Incremental.hs @@ -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 diff --git a/src/Futhark/Pass/Flatten/Intrablock.hs b/src/Futhark/Pass/Flatten/Intrablock.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/Flatten/Intrablock.hs @@ -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 diff --git a/src/Futhark/Pass/Flatten/Loop.hs b/src/Futhark/Pass/Flatten/Loop.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/Flatten/Loop.hs @@ -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 diff --git a/src/Futhark/Pass/Flatten/Match.hs b/src/Futhark/Pass/Flatten/Match.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/Flatten/Match.hs @@ -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 diff --git a/src/Futhark/Pass/Flatten/Monad.hs b/src/Futhark/Pass/Flatten/Monad.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/Flatten/Monad.hs @@ -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 () + } diff --git a/src/Futhark/Pass/Flatten/PreProcess.hs b/src/Futhark/Pass/Flatten/PreProcess.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/Flatten/PreProcess.hs @@ -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 diff --git a/src/Futhark/Pass/Flatten/SOAC.hs b/src/Futhark/Pass/Flatten/SOAC.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/Flatten/SOAC.hs @@ -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 diff --git a/src/Futhark/Pass/Flatten/WithAcc.hs b/src/Futhark/Pass/Flatten/WithAcc.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/Flatten/WithAcc.hs @@ -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 + } diff --git a/src/Futhark/Pass/NoGrid.hs b/src/Futhark/Pass/NoGrid.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/NoGrid.hs @@ -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 diff --git a/src/Futhark/Passes.hs b/src/Futhark/Passes.hs --- a/src/Futhark/Passes.hs +++ b/src/Futhark/Passes.hs @@ -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, diff --git a/src/Futhark/Tools.hs b/src/Futhark/Tools.hs --- a/src/Futhark/Tools.hs +++ b/src/Futhark/Tools.hs @@ -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 diff --git a/src/Futhark/Transform/FirstOrderTransform.hs b/src/Futhark/Transform/FirstOrderTransform.hs --- a/src/Futhark/Transform/FirstOrderTransform.hs +++ b/src/Futhark/Transform/FirstOrderTransform.hs @@ -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 diff --git a/src/Futhark/Transform/ISRWIM.hs b/src/Futhark/Transform/ISRWIM.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Transform/ISRWIM.hs @@ -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 diff --git a/src/Futhark/Transform/Rename.hs b/src/Futhark/Transform/Rename.hs --- a/src/Futhark/Transform/Rename.hs +++ b/src/Futhark/Transform/Rename.hs @@ -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 diff --git a/src/Futhark/Transform/Substitute.hs b/src/Futhark/Transform/Substitute.hs --- a/src/Futhark/Transform/Substitute.hs +++ b/src/Futhark/Transform/Substitute.hs @@ -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) diff --git a/src/Futhark/Transform/ToGPU.hs b/src/Futhark/Transform/ToGPU.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Transform/ToGPU.hs @@ -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 diff --git a/src/Futhark/Util.hs b/src/Futhark/Util.hs --- a/src/Futhark/Util.hs +++ b/src/Futhark/Util.hs @@ -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) diff --git a/src/Language/Futhark/Interpreter.hs b/src/Language/Futhark/Interpreter.hs --- a/src/Language/Futhark/Interpreter.hs +++ b/src/Language/Futhark/Interpreter.hs @@ -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 = diff --git a/src/Language/Futhark/Interpreter/Values.hs b/src/Language/Futhark/Interpreter/Values.hs --- a/src/Language/Futhark/Interpreter/Values.hs +++ b/src/Language/Futhark/Interpreter/Values.hs @@ -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 diff --git a/src/Language/Futhark/Pretty.hs b/src/Language/Futhark/Pretty.hs --- a/src/Language/Futhark/Pretty.hs +++ b/src/Language/Futhark/Pretty.hs @@ -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 diff --git a/src/Language/Futhark/Prop.hs b/src/Language/Futhark/Prop.hs --- a/src/Language/Futhark/Prop.hs +++ b/src/Language/Futhark/Prop.hs @@ -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] diff --git a/src/Language/Futhark/Traversals.hs b/src/Language/Futhark/Traversals.hs --- a/src/Language/Futhark/Traversals.hs +++ b/src/Language/Futhark/Traversals.hs @@ -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 diff --git a/src/Language/Futhark/TypeChecker.hs b/src/Language/Futhark/TypeChecker.hs --- a/src/Language/Futhark/TypeChecker.hs +++ b/src/Language/Futhark/TypeChecker.hs @@ -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 = diff --git a/src/Language/Futhark/TypeChecker/Constraints.hs b/src/Language/Futhark/TypeChecker/Constraints.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/TypeChecker/Constraints.hs @@ -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) diff --git a/src/Language/Futhark/TypeChecker/Consumption.hs b/src/Language/Futhark/TypeChecker/Consumption.hs --- a/src/Language/Futhark/TypeChecker/Consumption.hs +++ b/src/Language/Futhark/TypeChecker/Consumption.hs @@ -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] diff --git a/src/Language/Futhark/TypeChecker/Error.hs b/src/Language/Futhark/TypeChecker/Error.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/TypeChecker/Error.hs @@ -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) diff --git a/src/Language/Futhark/TypeChecker/Modules.hs b/src/Language/Futhark/TypeChecker/Modules.hs --- a/src/Language/Futhark/TypeChecker/Modules.hs +++ b/src/Language/Futhark/TypeChecker/Modules.hs @@ -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 diff --git a/src/Language/Futhark/TypeChecker/Monad.hs b/src/Language/Futhark/TypeChecker/Monad.hs --- a/src/Language/Futhark/TypeChecker/Monad.hs +++ b/src/Language/Futhark/TypeChecker/Monad.hs @@ -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 = diff --git a/src/Language/Futhark/TypeChecker/Names.hs b/src/Language/Futhark/TypeChecker/Names.hs --- a/src/Language/Futhark/TypeChecker/Names.hs +++ b/src/Language/Futhark/TypeChecker/Names.hs @@ -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 diff --git a/src/Language/Futhark/TypeChecker/Terms.hs b/src/Language/Futhark/TypeChecker/Terms.hs --- a/src/Language/Futhark/TypeChecker/Terms.hs +++ b/src/Language/Futhark/TypeChecker/Terms.hs @@ -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 "") (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 "") (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. diff --git a/src/Language/Futhark/TypeChecker/Terms/Loop.hs b/src/Language/Futhark/TypeChecker/Terms/Loop.hs --- a/src/Language/Futhark/TypeChecker/Terms/Loop.hs +++ b/src/Language/Futhark/TypeChecker/Terms/Loop.hs @@ -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. diff --git a/src/Language/Futhark/TypeChecker/Terms/Monad.hs b/src/Language/Futhark/TypeChecker/Terms/Monad.hs --- a/src/Language/Futhark/TypeChecker/Terms/Monad.hs +++ b/src/Language/Futhark/TypeChecker/Terms/Monad.hs @@ -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 diff --git a/src/Language/Futhark/TypeChecker/Terms/Pat.hs b/src/Language/Futhark/TypeChecker/Terms/Pat.hs --- a/src/Language/Futhark/TypeChecker/Terms/Pat.hs +++ b/src/Language/Futhark/TypeChecker/Terms/Pat.hs @@ -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 diff --git a/src/Language/Futhark/TypeChecker/Terms/Scope.hs b/src/Language/Futhark/TypeChecker/Terms/Scope.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/TypeChecker/Terms/Scope.hs @@ -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 diff --git a/src/Language/Futhark/TypeChecker/Terms/Unsized.hs b/src/Language/Futhark/TypeChecker/Terms/Unsized.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/TypeChecker/Terms/Unsized.hs @@ -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') diff --git a/src/Language/Futhark/TypeChecker/TySolve.hs b/src/Language/Futhark/TypeChecker/TySolve.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/TypeChecker/TySolve.hs @@ -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 #-} diff --git a/src/Language/Futhark/TypeChecker/Types.hs b/src/Language/Futhark/TypeChecker/Types.hs --- a/src/Language/Futhark/TypeChecker/Types.hs +++ b/src/Language/Futhark/TypeChecker/Types.hs @@ -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] -- diff --git a/src/Language/Futhark/TypeChecker/Unify.hs b/src/Language/Futhark/TypeChecker/Unify.hs --- a/src/Language/Futhark/TypeChecker/Unify.hs +++ b/src/Language/Futhark/TypeChecker/Unify.hs @@ -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' diff --git a/src/Language/Futhark/TypeChecker/UnionFind.hs b/src/Language/Futhark/TypeChecker/UnionFind.hs new file mode 100644 --- /dev/null +++ b/src/Language/Futhark/TypeChecker/UnionFind.hs @@ -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