futhark 0.21.10 → 0.21.11

raw patch · 114 files changed

+11151/−1964 lines, 114 filesdep +Diffdep ~futhark-datadep ~textPVP: major bump suggested

API removals or changes: PVP suggests a major version bump

Dependencies added: Diff

Dependency ranges changed: futhark-data, text

API changes (from Hackage documentation)

Files

• docs/conf.py +86/−7
• docs/error-index.rst +339/−92
• docs/installation.rst +18/−40
• docs/man/futhark-test.rst +7/−10
• futhark.cabal +24/−5
• prelude/ad.fut +23/−0
• prelude/array.fut +4/−4
• prelude/prelude.fut +1/−0
• prelude/soacs.fut +7/−7
• prelude/zip.fut +4/−4
• rts/c/chaselev.h +0/−176
• rts/c/cuda.h +3/−1
• rts/python/opencl.py +6/−1
• src/Futhark/AD/Derivatives.hs +386/−0
• src/Futhark/AD/Fwd.hs +486/−0
• src/Futhark/AD/Rev.hs +444/−0
• src/Futhark/AD/Rev/Loop.hs +456/−0
• src/Futhark/AD/Rev/Map.hs +233/−0
• src/Futhark/AD/Rev/Monad.hs +551/−0
• src/Futhark/AD/Rev/Reduce.hs +204/−0
• src/Futhark/AD/Rev/SOAC.hs +90/−0
• src/Futhark/AD/Rev/Scan.hs +154/−0
• src/Futhark/AD/Rev/Scatter.hs +188/−0
• src/Futhark/Analysis/LastUse.hs +4/−0
• src/Futhark/Analysis/Metrics.hs +2/−0
• src/Futhark/Analysis/PrimExp.hs +87/−2
• src/Futhark/Bench.hs +18/−17
• src/Futhark/CLI/Bench.hs +32/−43
• src/Futhark/CLI/Dev.hs +5/−0
• src/Futhark/CLI/LSP.hs +3/−3
• src/Futhark/CLI/Literate.hs +27/−1
• src/Futhark/CLI/Main.hs +144/−0
• src/Futhark/CLI/Test.hs +17/−14
• src/Futhark/CodeGen/Backends/CCUDA.hs +24/−2
• src/Futhark/CodeGen/Backends/COpenCL.hs +2/−2
• src/Futhark/CodeGen/Backends/GenericPython.hs +2/−2
• src/Futhark/CodeGen/ImpCode.hs +3/−1
• src/Futhark/CodeGen/ImpGen/GPU.hs +6/−0
• src/Futhark/CodeGen/OpenCL/Heuristics.hs +2/−2
• src/Futhark/CodeGen/RTS/C.hs +0/−6
• src/Futhark/Compiler.hs +10/−3
• src/Futhark/Compiler/CLI.hs +0/−7
• src/Futhark/Compiler/Config.hs +19/−2
• src/Futhark/Compiler/Program.hs +38/−19
• src/Futhark/Construct.hs +9/−0
• src/Futhark/Doc/Generator.hs +16/−15
• src/Futhark/IR/Aliases.hs +16/−3
• src/Futhark/IR/GPU/Op.hs +43/−2
• src/Futhark/IR/GPU/Simplify.hs +35/−2
• src/Futhark/IR/Mem/Simplify.hs +1/−1
• src/Futhark/IR/Parse.hs +16/−1
• src/Futhark/IR/Primitive.hs +1/−1
• src/Futhark/IR/Prop/Aliases.hs +1/−1
• src/Futhark/IR/Prop/Names.hs +6/−0
• src/Futhark/IR/Prop/Types.hs +8/−2
• src/Futhark/IR/SOACS.hs +25/−0
• src/Futhark/IR/SOACS/SOAC.hs +48/−2
• src/Futhark/IR/SOACS/Simplify.hs +22/−7
• src/Futhark/IR/SegOp.hs +19/−8
• src/Futhark/IR/Syntax.hs +15/−6
• src/Futhark/IR/Syntax/Core.hs +5/−0
• src/Futhark/Internalise/Defunctorise.hs +9/−1
• src/Futhark/Internalise/Exps.hs +14/−0
• src/Futhark/LSP/Compile.hs +54/−28
• src/Futhark/LSP/Handlers.hs +15/−16
• src/Futhark/LSP/PositionMapping.hs +80/−0
• src/Futhark/LSP/State.hs +39/−3
• src/Futhark/LSP/Tool.hs +57/−5
• src/Futhark/Optimise/BlkRegTiling.hs +1199/−903
• src/Futhark/Optimise/CSE.hs +11/−4
• src/Futhark/Optimise/Fusion.hs +5/−4
• src/Futhark/Optimise/GenRedOpt.hs +432/−0
• src/Futhark/Optimise/HistAccs.hs +163/−0
• src/Futhark/Optimise/MergeGPUBodies.hs +634/−0
• src/Futhark/Optimise/ReduceDeviceSyncs.hs +848/−0
• src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs +1620/−0
• src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable/Graph.hs +569/−0
• src/Futhark/Optimise/Simplify.hs +1/−1
• src/Futhark/Optimise/Simplify/Engine.hs +38/−4
• src/Futhark/Optimise/Simplify/Rep.hs +6/−5
• src/Futhark/Optimise/Sink.hs +50/−1
• src/Futhark/Optimise/TileLoops.hs +33/−26
• src/Futhark/Optimise/TileLoops/Shared.hs +140/−17
• src/Futhark/Pass.hs +3/−3
• src/Futhark/Pass/AD.hs +102/−0
• src/Futhark/Pass/ExpandAllocations.hs +1/−0
• src/Futhark/Pass/ExplicitAllocations/GPU.hs +2/−0
• src/Futhark/Pass/ExtractMulticore.hs +4/−0
• src/Futhark/Passes.hs +35/−3
• src/Futhark/Pipeline.hs +11/−10
• src/Futhark/Script.hs +2/−2
• src/Futhark/Test.hs +6/−2
• src/Futhark/Transform/FirstOrderTransform.hs +4/−0
• src/Futhark/Util.hs +20/−8
• src/Futhark/Util/Console.hs +5/−0
• src/Futhark/Util/IntegralExp.hs +2/−0
• src/Futhark/Util/ProgressBar.hs +57/−0
• src/Language/Futhark/Interpreter.hs +4/−0
• src/Language/Futhark/Prop.hs +82/−26
• src/Language/Futhark/Query.hs +1/−4
• src/Language/Futhark/Syntax.hs +64/−56
• src/Language/Futhark/TypeChecker/Modules.hs +48/−49
• src/Language/Futhark/TypeChecker/Monad.hs +2/−1
• src/Language/Futhark/TypeChecker/Terms.hs +7/−3
• src/Language/Futhark/TypeChecker/Terms/Monad.hs +13/−7
• src/Language/Futhark/TypeChecker/Terms/Pat.hs +6/−3
• src/Language/Futhark/TypeChecker/Types.hs +35/−25
• src/Language/Futhark/TypeChecker/Unify.hs +31/−14
• src/Language/Futhark/Warnings.hs +1/−1
• src/main.hs +4/−140
• unittests/Futhark/AD/DerivativesTests.hs +50/−0
• unittests/Futhark/IR/Mem/IxFunTests.hs +1/−0
• unittests/Language/Futhark/TypeChecker/TypesTests.hs +84/−65
• unittests/futhark_tests.hs +2/−0

@@ -16,8 +16,11 @@ import sys import os import re-from pygments.lexer import RegexLexer+from pygments.lexer import RegexLexer, bygroups from pygments import token+from pygments import unistring as uni+from pygments.token import Text, Comment, Operator, Keyword, Name, String, \+    Number, Punctuation, Whitespace from sphinx.highlighting import lexers  # If extensions (or modules to document with autodoc) are in another directory,@@ -103,18 +106,94 @@ # The name of the Pygments (syntax highlighting) style to use. pygments_style = 'sphinx' + class FutharkLexer(RegexLexer):+    """+    A Futhark lexer++    .. versionadded:: 2.8+    """     name = 'Futhark'+    url = 'https://futhark-lang.org/'+    aliases = ['futhark']+    filenames = ['*.fut']+    mimetypes = ['text/x-futhark'] +    num_types = ('i8', 'i16', 'i32', 'i64', 'u8', 'u16', 'u32', 'u64', 'f32', 'f64')++    other_types = ('bool', )++    reserved = ('if', 'then', 'else', 'def', 'let', 'loop', 'in', 'with', 'type',+                'val', 'entry', 'for', 'while', 'do', 'case', 'match',+                'include', 'import', 'module', 'open', 'local', 'assert', '_')++    ascii = ('NUL', 'SOH', '[SE]TX', 'EOT', 'ENQ', 'ACK',+             'BEL', 'BS', 'HT', 'LF', 'VT', 'FF', 'CR', 'S[OI]', 'DLE',+             'DC[1-4]', 'NAK', 'SYN', 'ETB', 'CAN',+             'EM', 'SUB', 'ESC', '[FGRU]S', 'SP', 'DEL')++    num_postfix = r'(%s)?' % '|'.join(num_types)++    identifier_re = '[a-zA-Z_][a-zA-Z_0-9\']*'++    # opstart_re = '+\-\*/%=\!><\|&\^'+     tokens = {         'root': [-            (r'(if|then|else|let|loop|in|val|for|do|with|local|open|include|import|type|def|entry|module|while|module)\b', token.Keyword),-            (r"#?[a-zA-Z_][a-zA-Z0-9_']*", token.Name),-            (r"--.*", token.Comment),-            (r'.', token.Text)-        ]-    }+            (r'--(.*?)$', Comment.Single),+            (r'\s+', Whitespace),+            (r'\(\)', Punctuation),+            (r'\b(%s)(?!\')\b' % '|'.join(reserved), Keyword.Reserved),+            (r'\b(%s)(?!\')\b' % '|'.join(num_types + other_types), Keyword.Type), +            # Identifiers+            (r'#\[([a-zA-Z_\(\) ]*)\]', Comment.Preproc),+            (r'[#!]?(%s\.)*%s' % (identifier_re, identifier_re), Name),++            (r'\\', Operator),+            (r'[-+/%=!><|&*^][-+/%=!><|&*^.]*', Operator),+            (r'[][(),:;`{}?.\']', Punctuation),++            #  Numbers+            (r'0[xX]_*[\da-fA-F](_*[\da-fA-F])*_*[pP][+-]?\d(_*\d)*' + num_postfix,+             Number.Float),+            (r'0[xX]_*[\da-fA-F](_*[\da-fA-F])*\.[\da-fA-F](_*[\da-fA-F])*'+             r'(_*[pP][+-]?\d(_*\d)*)?' + num_postfix, Number.Float),+            (r'\d(_*\d)*_*[eE][+-]?\d(_*\d)*' + num_postfix, Number.Float),+            (r'\d(_*\d)*\.\d(_*\d)*(_*[eE][+-]?\d(_*\d)*)?' + num_postfix, Number.Float),+            (r'0[bB]_*[01](_*[01])*' + num_postfix, Number.Bin),+            (r'0[xX]_*[\da-fA-F](_*[\da-fA-F])*' + num_postfix, Number.Hex),+            (r'\d(_*\d)*' + num_postfix, Number.Integer),++            #  Character/String Literals+            (r"'", String.Char, 'character'),+            (r'"', String, 'string'),+            #  Special+            (r'\[[a-zA-Z_\d]*\]', Keyword.Type),+            (r'\(\)', Name.Builtin),+        ],+        'character': [+            # Allows multi-chars, incorrectly.+            (r"[^\\']'", String.Char, '#pop'),+            (r"\\", String.Escape, 'escape'),+            ("'", String.Char, '#pop'),+        ],+        'string': [+            (r'[^\\"]+', String),+            (r"\\", String.Escape, 'escape'),+            ('"', String, '#pop'),+        ],++        'escape': [+            (r'[abfnrtv"\'&\\]', String.Escape, '#pop'),+            (r'\^[][' + uni.Lu + r'@^_]', String.Escape, '#pop'),+            ('|'.join(ascii), String.Escape, '#pop'),+            (r'o[0-7]+', String.Escape, '#pop'),+            (r'x[\da-fA-F]+', String.Escape, '#pop'),+            (r'\d+', String.Escape, '#pop'),+            (r'(\s+)(\\)', bygroups(Whitespace, String.Escape), '#pop'),+        ],+    }  lexers['futhark'] = FutharkLexer() 

@@ -16,8 +16,10 @@  A core principle of uniqueness typing (see :ref:`in-place-updates`) is that after a variable is "consumed", it must not be used again.  For-example, this is invalid, and will result in the error above::+example, this is invalid, and will result in the error above: +.. code-block:: futhark+   let y = x with [0] = 0   in x @@ -27,14 +29,18 @@ consumed, its *aliases* are also considered consumed.  Aliasing is the possibility of two variables occupying the same memory at run-time. For example, this will fail as above, because ``y`` and ``x`` are-aliased::+aliased: +.. code-block:: futhark+   let y = x   let z = y with [0] = 0   in x -We can always break aliasing by using a ``copy`` expression::+We can always break aliasing by using a ``copy`` expression: +.. code-block:: futhark+   let y = copy x   let z = y with [0] = 0   in x@@ -46,17 +52,21 @@  This error message occurs for programs that try to perform a consumption (such as an in-place update) on variables that are not-consumable.  For example, it would occur for the following program::+consumable.  For example, it would occur for the following program: -  let f (a: []i32) =+.. code-block:: futhark++  def f (a: []i32) =     let a[0] = a[0]+1     in a  Only arrays with a a *unique array type* can be consumed.  Such a type is written by prefixing the array type with an asterisk.  The program-could be fixed by writing it like this::+could be fixed by writing it like this: -  let f (a: *[]i32) =+.. code-block:: futhark++  def f (a: *[]i32) =     let a[0] = a[0]+1     in a @@ -65,9 +75,11 @@ :ref:`in-place-updates` for the full details.  You can always obtain a unique copy of an array by using-``copy``::+``copy``: -  let f (a: []i32) =+.. code-block:: futhark++  def f (a: []i32) =     let a = copy a     let a[0] = a[0]+1     in a@@ -80,18 +92,22 @@ "Unique-typed return value of *x* is aliased to *y*, which is not consumable" ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -This can be caused by a function like this::+This can be caused by a function like this: -  let f (xs: []i32) : *[]i32 = xs+.. code-block:: futhark +  def f (xs: []i32) : *[]i32 = xs+ We are saying that ``f`` returns a *unique* array - meaning it has no aliases - but at the same time, it aliases the parameter *xs*, which is not marked as being unique (see :ref:`in-place-updates`).  This violates one of the core guarantees provided by uniqueness types, namely that a unique return value does not alias any value that might be used in the future.  Imagine if this was permitted, and we had a-program that used ``f``::+program that used ``f``: +.. code-block:: futhark+   let b = f a   let b[0] = x   ...@@ -102,10 +118,12 @@  As with most uniqueness errors, it can be fixed by using ``copy xs`` to break the aliasing.  We can also change the type of ``f`` to take a-unique array as input::+unique array as input: -  let f (xs: *[]i32) : *[]i32 = xs+.. code-block:: futhark +  def f (xs: *[]i32) : *[]i32 = xs+ This makes ``xs`` "consumable", in the sense used by the error message.  .. _unique-return-aliased:@@ -113,46 +131,55 @@ "A unique-typed component of the return value of *x* is aliased to some other component" ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -Caused by programs like the following::+Caused by programs like the following: -  let main (xs: *[]i32) : (*[]i32, *[]i32) = (xs, xs)+.. code-block:: futhark +  def main (xs: *[]i32) : (*[]i32, *[]i32) = (xs, xs)+ While we are allowed to "consume" ``xs``, as it is a unique parameter, this function is trying to return two unique values that alias each other.  This violates one of the core guarantees provided by uniqueness types, namely that a unique return value does not alias any value that might be used in the future (see :ref:`in-place-updates`) - and in this case, the two values alias each other.  We can fix this by-inserting copies to break the aliasing::+inserting copies to break the aliasing: -  let main (xs: *[]i32) : (*[]i32, *[]i32) = (xs, copy xs)+.. code-block:: futhark +  def main (xs: *[]i32) : (*[]i32, *[]i32) = (xs, copy xs)+ .. _consuming-parameter:  "Consuming parameter passed non-unique argument" ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -Caused by programs like the following::+Caused by programs like the following: -  let update (xs: *[]i32) = xs with [0] = 0+.. code-block:: futhark -  let f (ys: []i32) = update ys+  def update (xs: *[]i32) = xs with [0] = 0 +  def f (ys: []i32) = update ys+ The update ``function`` *consumes* its ``xs`` argument to perform an :ref:`in-place update <in-place-updates>`, as denoted by the asterisk before the type.  However, the ``f`` function tries to pass an array that it is not allowed to consume (no asterisk before the type). - One solution is to change the type of ``f`` so that it also consumes-its input, which allows it to pass it on to ``update``::+its input, which allows it to pass it on to ``update``: -  let f (ys: *[]i32) = update ys+.. code-block:: futhark -Another solution to ``copy`` the array that we pass to ``update``::+  def f (ys: *[]i32) = update ys -  let f (ys: []i32) = update (copy 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."@@ -160,75 +187,97 @@  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-that does::+that does: -  let apply 'a 'b (f: a -> b) (x: a) = f x+.. code-block:: futhark -  let consume (xs: *[]i32) = xs with [0] = 0+  def apply 'a 'b (f: a -> b) (x: a) = f x -  let f (arr: *[]i32) = apply consume arr+  def consume (xs: *[]i32) = xs with [0] = 0 +  def f (arr: *[]i32) = apply consume arr+ We can fix this by changing ``consume`` so that it does not have to-consume its argument, by adding a ``copy``::+consume its argument, by adding a ``copy``: -  let consume (xs: []i32) = copy xs with [0] = 0+.. code-block:: futhark +  def consume (xs: []i32) = copy xs with [0] = 0+ Or we can create a variant of ``apply`` that accepts a consuming-function::+function: -  let apply 'a 'b (f: *a -> b) (x: *a) = f x+.. code-block:: futhark +  def apply 'a 'b (f: *a -> b) (x: *a) = f x+ .. _alias-free-variable:  "Function result aliases the free variable *x*" ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -Caused by definitions such as the following::+Caused by definitions such as the following: -  let x = [1,2,3]+.. code-block:: futhark -  let f () = x+  def x = [1,2,3] +  def f () = x+ To simplify the tracking of aliases, the Futhark type system requires that the result of a function may only alias the function parameters,-not any free variables.  Use ``copy`` to fix this::+not any free variables.  Use ``copy`` to fix this: -  let f () = copy x+.. code-block:: futhark +  def f () = copy x+ .. _inaccessible-size:  "Parameter *x* refers to size *y* which will not be accessible to the caller ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~  This happens when the size of an array parameter depends on a name-that cannot be expressed in the function type::+that cannot be expressed in the function type: -  let f (x: i64, y: i64) (A: [x]bool) = true+.. code-block:: futhark -Intuitively, this function might have the following type::+  def f (x: i64, y: i64) (A: [x]bool) = true +Intuitively, this function might have the following type:++.. code-block:: futhark+   val f : (x: i64, y: i64) -> [x]bool -> bool  But this is not currently a valid Futhark type.  In a function type, each parameter can be named *as a whole*, but it cannot be taken apart in a pattern.  In this case, we could fix it by splitting the tuple-parameter into two separate parameters::+parameter into two separate parameters: -  let f (x: i64) (y: i64) (A: [x]bool) = true+.. code-block:: futhark -This gives the following type::+  def f (x: i64) (y: i64) (A: [x]bool) = true +This gives the following type:++.. code-block:: futhark+   val f : (x: i64) -> (y: i64) -> [x]bool -> bool  Another workaround is to loosen the static safety, and use a size-coercion to give A its expected size::+coercion to give A its expected size: -  let f (x: i64, y: i64) (A_unsized: []bool) =+.. code-block:: futhark++  def f (x: i64, y: i64) (A_unsized: []bool) =     let A = A_unsized :> [x]bool     in true -This will produce a function with the following type::+This will produce a function with the following type: +.. code-block:: futhark+   val f [d] : (i64, i64) -> [d]bool -> bool  This does however lose the constraint that the size of the array must@@ -236,10 +285,12 @@ fail at run-time.  The error is not always due to an explicit type annotation.  It might-also be due to size inference::+also be due to size inference: -  let f (x: i64, y: i64) (A: []bool) = zip A (iota x)+.. code-block:: futhark +  def f (x: i64, y: i64) (A: []bool) = zip A (iota x)+ Here the type rules force ``A`` to have size ``x``, leading to a problematic type.  It can be fixed using the techniques above. @@ -251,14 +302,18 @@ "Size *x* unused in pattern." ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -Caused by expressions like this::+Caused by expressions like this: -  let [n] (y: i32) = x+.. code-block:: futhark -And functions like this::+  def [n] (y: i32) = x -  let f [n] (x: i32) = x+And functions like this: +.. code-block:: futhark++  def f [n] (x: i32) = x+ Since ``n`` is not the size of anything, it cannot be assigned a value at runtime.  Hence this program is rejected. @@ -269,9 +324,11 @@  Causality check errors occur when the program is written in such a way that a size is needed before it is actually computed.  See-:ref:`causality` for the full rules.  Contrived example::+:ref:`causality` for the full rules.  Contrived example: -  let f (b: bool) (xs: []i32) =+.. code-block:: futhark++  def f (b: bool) (xs: []i32) =     let a = [] : [][]i32     let b = [filter (>0) xs]     in a[0] == b[0]@@ -293,9 +350,11 @@  This error occurs when you define a function that can never be applied, as it requires an input of a specific size, and that size is-not known.  Somewhat contrived example::+not known.  Somewhat contrived example: -  let f (x: bool) =+.. code-block:: futhark++  def f (x: bool) =     let n = if x then 10 else 20     in \(y: [n]bool) -> ... @@ -311,9 +370,11 @@ actually mean to.  However, in the case that that the above really is what you intend, the workaround is to make the function fully polymorphic, and then perform a size coercion to the desired size-inside the function body itself::+inside the function body itself: -  let f (x: bool) =+.. code-block:: futhark++  def f (x: bool) =     let n = if x then 10 else 20     in \(y_any: []bool) ->          let y = y_any :> [n]bool@@ -328,17 +389,23 @@ ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~  This occurs most commonly when we use function composition with one or-more functions that return an *existential size*.  Example::+more functions that return an *existential size*.  Example: +.. code-block:: futhark+   filter (>0) >-> length -The ``filter`` function has this type::+The ``filter`` function has this type: +.. code-block:: futhark+   val filter [n] 't : (t -> bool) -> [n]t -> ?[m].[m]t  That is, ``filter`` returns an array whose size is not known until the-function actually returns.  The ``length`` function has this type::+function actually returns.  The ``length`` function has this type: +.. code-block:: futhark+   val length [n] 't : [n]t -> i64  Whenever ``length`` occurs (as in the composition above), the type@@ -351,8 +418,10 @@ type of ``length``, and therefore the program is rejected.  The common workaround is to use *pipelining* instead of composition-whenever we use functions with existential return types::+whenever we use functions with existential return types: +.. code-block:: futhark+   xs |> filter (>0) |> length  This works because ``|>`` is left-associative, and hence the ``xs |>@@ -362,6 +431,155 @@ We can of course also write it as ``length (filter (>0) xs)``, with no use of either pipelining or composition. +.. _unused-existential:++"Existential size *n* not used as array size"+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~++This error occurs for type expressions that use explicit existential+quantification in an incorrect way, such as the following examples:++.. code-block:: futhark++  ?[n].bool++  ?[n].bool -> [n]bool++When we use existential quantification, we are required to use the+size within its scope, *and* it must not exclusively be used to the+right of function arrow.++To understand the motivation behind this rule, consider that when we+use an existential quantifier we are saying that there is *some size*,+it just cannot be known statically, but must be read from some value+(i.e. array) at runtime.  In the first example above, the existential+size ``n`` is not used at all, so the actual value cannot be+determined at runtime.  In the second example, while an array+``[n]bool`` does exist, it is part of a function type, and at runtime+functions are black boxes and don't "carry" the size of their+parameter or result types.++The workaround is to actually use the existential size.  This can be+as simple as adding a *witness array* of type ``[n]()``:++.. code-block:: futhark++  ?[n].([n](),bool)++  ?[n].([n](), bool -> [n]bool)++Such an array will take up no space at runtime.++.. _unify-param-existential:++"Parameter *x* used as size would go out of scope."+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~++This error tends to happen because higher-order functions are used in+a way that causes a size requirement to become impossible to+constrait.  Real programs that run into this issue are quite complex,+but to illustrate the problem, consider the following contrived+function:++.. code-block:: futhark++  def f (n: i64) (m: i64) (b: [n][m]bool) = b[0,0]++We have the following type:++.. code-block:: futhark++  val f : (n: i64) -> (m: i64) -> (b: [n][m]bool) -> bool++Now suppose we say:++.. code-block:: futhark++  def g = uncurry f++What should be the type of ``g``?  Intuitively, something like this:++.. code-block:: futhark++  val g : (n: i64, m: i64) -> (b: [n][m]bool) -> bool++But this is *not* expressible in the Futhark type system - and even if+it were, it would not be easy to infer this in general, as it depends+on exactly what ``uncurry`` does, which the type checker does not+know.++As a workaround, we can use explicit type annotation and size+coercions to give ``g`` an acceptable type:++.. code-block:: futhark++  def g [a][b] (n,m) (b: [a][b]bool) = f n m (b :> [n][m]bool)++Another workaround, which is often the right one in cases not as+contrived as above, is to modify ``f`` itself to produce a *witness*+of the constraint, in the form of an array of shape ``[n][m]``:++.. code-block:: futhark++  def f (n: i64) (m: i64) : ([n][m](), [n][m]bool -> bool) =+    (replicate n (replicate m ()), \b -> b[0,0])++Then ``uncurry f`` works just fine and has the following type:++.. code-block:: futhark++  (i64, i64) -> ?[n][m].([n][m](), [n][m]bool -> bool)++Programming with such *explicit size witnesses* is a fairly advanced+technique, but often necessary when writing advanced size-dependent+code.++.. _ambiguous-size:++"Ambiguous size *x*"+~~~~~~~~~~~~~~~~~~~~++There are various sources for this error, but they all have the same+ultimate cause: the type checker cannot figure out how some symbolic+size name should be resolved to a concrete size.  The simplest+example, although contrived, is probably this:++.. code-block:: futhark++   let [n][m] (xss: [n][m]i64) = []++The type checker can infer that ``n`` should be zero, but how can it+possibly figure out the shape of the (non-existent) rows of the+two-dimensional array?  This can be fixed in many ways, but adding a+type ascription to the array is one of them: ``[] : [0][2]i64``.++Another common case arises when using holes.  For an expression+``length ???``, how would the type checker figure out the intended+size of the array that the hole represents?  Again, this can be solved+with a type ascription: ``length (??? : [10]bool)``.++Finally, ambiguous sizes can also occur for functions that use size+parameters only in "non-witnessing" position, meaning sizes that are+not actually uses as sizes of real arrays.  An example:++.. code-block:: futhark++   def f [n] (g: [n]i64 -> i64) : i64 = n++   def main = f (\xs -> xs[0])++Note that ``f`` is a higher order function, and that the size+parameter ``n`` is only used in the type of the ``g`` function.+Futhark's value model is such that given a value of type ``[n]i64 ->+i64``, we cannot extract an ``n`` from it.  Using a function such as+``f`` is only valid when ``n`` can be inferred from the usage, which+is not the case here.  Again, we can fix it by adding a type+ascription to disambiguate:++.. code-block:: futhark++   def main = f (\(xs:[1]i64) -> xs[0])+ Module errors ------------- @@ -371,18 +589,22 @@ ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~  This occurs when the program uses the ``entry`` keyword inside a-module::+module: +.. code-block:: futhark+   module m = {     entry f x = x + 1   }  Entry points can only be declared at the top level of a file.  When we wish to make a function from inside a module available as an entry-point, we must define a wrapper function::+point, we must define a wrapper function: +.. code-block:: futhark+   module m = {-    let f x = x + 1+    def f x = x + 1   }    entry f = m.f@@ -390,18 +612,22 @@ .. _module-is-parametric:  "Module *x* is a parametric module-----------------------------------+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -A parametric module is a module-level function::+A parametric module is a module-level function: +.. code-block:: futhark+   module PM (P: {val x : i64}) = {-    let y = x + 2+    def y = x + 2   }  If we directly try to access the component of ``PM``, as ``PM.y``, we will get an error.  To use ``PM`` we must first apply it to a module-of the expected type::+of the expected type: +.. code-block:: futhark+   module M = PM { val x = 2 : i64 }  Now we can say ``M.y``.  See :ref:`module-system` for more.@@ -416,14 +642,18 @@  This occurs for overloaded constants such as ``1234`` that are inferred by context to have a type that is too narrow for their value.-Example::+Example: +.. code-block::+   257 : u8  It is not an error to have a *non-overloaded* numeric constant whose value is too large for its type.  The following is perfectly-cromulent::+cromulent: +.. code-block::+   257u8  In such cases, the behaviour is overflow (so this is equivalent to@@ -435,30 +665,38 @@ ~~~~~~~~~~~~~~~~~~~  There are various cases where the type checker is unable to infer the-full type of something.  For example::+full type of something.  For example: -  let f r = r.x+.. code-block:: futhark +  def f r = r.x+ We know that ``r`` must be a record with a field called ``x``, but maybe the record could also have other fields as well.  Instead of assuming a perhaps too narrow type, the type checker signals an error. The solution is always to add a type annotation in one or more places-to disambiguate the type::+to disambiguate the type: -  let f (r: {x:bool, y:i32}) = r.x+.. code-block:: futhark +  def f (r: {x:bool, y:i32}) = r.x+ Usually the best spot to add such an annotation is on a function parameter, as above.  But for ambiguous sum types, we often have to-put it on the return type.  Consider::+put it on the return type.  Consider: -  let f (x: bool) = #some x+.. code-block:: futhark +  def f (x: bool) = #some x+ The type of this function is ambiguous, because the type checker must know what other possible contructors (apart from ``#some``) are-possible.  We fix it with a type annotation on the return type::+possible.  We fix it with a type annotation on the return type: -  let f (x: bool) : (#some bool | #none) = #just x+.. code-block:: futhark +  def f (x: bool) : (#some bool | #none) = #just x+ See :ref:`typeabbrevs` for how to avoid typing long types in several places. @@ -477,26 +715,31 @@ ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~  Futhark requires ``match`` expressions to be *exhaustive* - that is,-cover all possible forms of the value being pattern-matches.-Example::+cover all possible forms of the value being matched.  Example: -  let f (x: i32) =+.. code-block:: futhark++  def f (x: i32) =     match x case 0 -> false             case 1 -> true  Usually this is an actual bug, and you fix it by adding the missing cases.  But sometimes you *know* that the missing cases will never actually occur at run-time.  To satisfy the type checker, you can turn-the final case into a wildcard that matches anything::+the final case into a wildcard that matches anything: -  let f (x: i32) =+.. code-block:: futhark++  def f (x: i32) =     match x case 0 -> false             case _ -> true  Alternatively, you can add a wildcard case that explicitly asserts-that it should never happen::+that it should never happen: -  let f (x: i32) =+.. code-block:: futhark++  def f (x: i32) =     match x case 0 -> false             case 1 -> true             case _ -> assert false false@@ -511,13 +754,17 @@ When performing a :ref:`record update <record_update>`, the type of the field we are updating must be known.  This restriction is based on a limitation in the type type checker, so the notion of "known" is a bit-subtle::+subtle: -  let f r : {x:i32} = r with x = 0+.. code-block:: futhark +  def f r : {x:i32} = r with x = 0+ Even though the return type annotation disambiguates the type, this program still fails to type check.  This is because the return type is not consulted until *after* the body of the function has been checked.-The solution is to put a type annotation on the parameter instead::+The solution is to put a type annotation on the parameter instead: -  let f (r : {x:i32}) = r with x = 0+.. code-block:: futhark++  def f (r : {x:i32}) = r with x = 0

@@ -36,14 +36,13 @@ Compiling from source --------------------- -The recommended way to compile Futhark is with the `Haskell Tool-Stack`_, which handles dependencies and compilation of the Futhark-compiler.  You will therefore need to install the ``stack`` tool.-Fortunately, the ``stack`` developers provide ample documentation-about `installing Stack`_ on a multitude of operating systems.  If-you're lucky, it may even be in your local package repository.+To compile Futhark you must first install an appropriate version of+GHC, either with [ghcup](https://www.haskell.org/ghcup/) or a package+manager.  Any version since GHC 8.10 should work.  You also need the+``cabal`` command line program, which ghcup will install for you as+well. -You can either retrieve a `source release tarball+You then either retrieve a `source release tarball <https://github.com/diku-dk/futhark/releases>`_ or perform a checkout of our Git repository:: @@ -53,48 +52,27 @@    $ cd futhark -To get all the prerequisites for building the Futhark compiler-(including, if necessary, the appropriate version of the Haskell-compiler), run::+First you must run the following command to download metadata about+Futhark's dependencies:: -  $ stack setup+  $ make configure -Note that this will not install anything system-wide and will have no-effect outside the Futhark build directory. Now you can run the-following command to build the Futhark compiler, including all-dependencies::+To build the Futhark compiler and all of its dependencies, run:: -  $ stack build+  $ make build -The Futhark compiler and its tools will now be built.  This step-typically requires at least 8GiB of memory.  You may be able to build-it on a smaller machine by adding the ``--fast`` option, although the-resulting Futhark compiler binary will run slower.+This step typically requires at least 8GiB of memory.  This will+create files in your ``~/.cabal`` and ``~/.ghc`` directories.  After building, you can copy the binaries to your ``$HOME/.local/bin`` directory by running:: -  $ stack install--Note that this does not install the Futhark manual pages.--Compiling with ``cabal``-~~~~~~~~~~~~~~~~~~~~~~~~--You can also compile Futhark with ``cabal``.  If so, you must install-an appropriate version of GHC (usually the newest) and ``cabal``-yourself, for example through your favourite package manager.  On-Linux, you can always use `ghcup-<https://gitlab.haskell.org/haskell/ghcup>`_.  Then clone the-repository as listed above and run::--  $ cabal update-  $ cabal build--To install the Futhark binaries to a specific location, for example-``$HOME/.local/bin``, run::+  $ make install -  $ cabal install --install-method=copy  --overwrite-policy=always --installdir=$HOME/.local/bin/+You can set the ``PREFIX`` environment variable to indicate a+different installation path.  Note that this does not install the+Futhark manual pages.  You can delete ``~/.cabal`` and ``~/.ghc``+after this if you wish - the ``futhark`` binary will still work.  Installing from a precompiled snapshot --------------------------------------

@@ -14,9 +14,10 @@ DESCRIPTION =========== -This tool is used to test Futhark programs based on input/output-datasets.  If a directory is given, all contained files with a-``.fut`` extension are considered.+Test Futhark programs based on input/output datasets.  All contained+``.fut`` files within a given directory are considered.  By default,+tests are carried out with compiled code.  This can be changed with+the ``-i`` option.  A Futhark test program is an ordinary Futhark program, with at least one test block describing input/output test cases and possibly other@@ -89,10 +90,6 @@  This is used to test the type checker. -By default, both the interpreter and compiler is run on all test cases-(except those that have specified ``compiled``), although this can be-changed with command-line options to ``futhark test``.- Tuple syntax is not supported when specifying input and output values. Instead, you can write an N-tuple as its constituent N values.  Beware of syntax errors in the values - the errors reported by@@ -135,12 +132,12 @@   By default, ``--cache-file`` is not passed.  -c-  Only run compiled code - do not run any interpreters.+  Only run compiled code - do not run the interpreter.  This is the+  default.  -C   Compile the programs, but do not run them. - --concurrency=NUM    The number of tests to run concurrently.  Defaults to the number of@@ -153,7 +150,7 @@   where *foo* is the backend used.  -i-  Only interpret - do not run any compilers.+  Test with the interpreter.  -t   Type-check the programs, but do not run them.

@@ -1,6 +1,6 @@ cabal-version: 2.4 name:           futhark-version:        0.21.10+version:        0.21.11 synopsis:       An optimising compiler for a functional, array-oriented language.  description:    Futhark is a small programming language designed to be compiled to@@ -55,6 +55,16 @@   exposed-modules:       Futhark       Futhark.Actions+      Futhark.AD.Derivatives+      Futhark.AD.Fwd+      Futhark.AD.Rev+      Futhark.AD.Rev.Loop+      Futhark.AD.Rev.Map+      Futhark.AD.Rev.Monad+      Futhark.AD.Rev.Reduce+      Futhark.AD.Rev.Scan+      Futhark.AD.Rev.Scatter+      Futhark.AD.Rev.SOAC       Futhark.Analysis.Alias       Futhark.Analysis.CallGraph       Futhark.Analysis.DataDependencies@@ -88,6 +98,7 @@       Futhark.CLI.Doc       Futhark.CLI.Literate       Futhark.CLI.LSP+      Futhark.CLI.Main       Futhark.CLI.Misc       Futhark.CLI.Multicore       Futhark.CLI.MulticoreWASM@@ -216,6 +227,7 @@       Futhark.LSP.Handlers       Futhark.LSP.Tool       Futhark.LSP.State+      Futhark.LSP.PositionMapping       Futhark.MonadFreshNames       Futhark.Optimise.BlkRegTiling       Futhark.Optimise.CSE@@ -223,12 +235,18 @@       Futhark.Optimise.Fusion       Futhark.Optimise.Fusion.Composing       Futhark.Optimise.Fusion.LoopKernel+      Futhark.Optimise.GenRedOpt+      Futhark.Optimise.HistAccs       Futhark.Optimise.InPlaceLowering       Futhark.Optimise.InPlaceLowering.LowerIntoStm       Futhark.Optimise.InPlaceLowering.SubstituteIndices       Futhark.Optimise.InliningDeadFun       Futhark.Optimise.MemoryBlockMerging       Futhark.Optimise.MemoryBlockMerging.GreedyColoring+      Futhark.Optimise.MergeGPUBodies+      Futhark.Optimise.ReduceDeviceSyncs+      Futhark.Optimise.ReduceDeviceSyncs.MigrationTable+      Futhark.Optimise.ReduceDeviceSyncs.MigrationTable.Graph       Futhark.Optimise.Simplify       Futhark.Optimise.Simplify.Engine       Futhark.Optimise.Simplify.Rep@@ -244,6 +262,7 @@       Futhark.Optimise.TileLoops.Shared       Futhark.Optimise.Unstream       Futhark.Pass+      Futhark.Pass.AD       Futhark.Pass.ExpandAllocations       Futhark.Pass.ExplicitAllocations       Futhark.Pass.ExplicitAllocations.GPU@@ -284,6 +303,7 @@       Futhark.Util.Log       Futhark.Util.Options       Futhark.Util.Pretty+      Futhark.Util.ProgressBar       Futhark.Util.Table       Futhark.Version       Language.Futhark@@ -336,13 +356,14 @@     , bmp >=1.2.6.3     , containers >=0.6.2.1     , cryptohash-md5+    , Diff >=0.4.1     , directory >=1.3.0.0     , directory-tree >=0.12.1     , dlist >=0.6.0.1     , file-embed >=0.0.14.0     , filepath >=1.4.1.1     , free >=4.12.4-    , futhark-data >= 1.0.3.0+    , futhark-data >= 1.1.0.0     , futhark-server >= 1.1.2.1     , futhark-manifest >= 1.0.0.0     , githash >=0.1.6.1@@ -379,19 +400,17 @@  executable futhark   main-is: src/main.hs-  other-modules:-      Paths_futhark   ghc-options: -Wall -Wcompat -Wno-incomplete-uni-patterns -Wredundant-constraints -Wincomplete-record-updates -Wmissing-export-lists -Wunused-packages -threaded -rtsopts "-with-rtsopts=-N -qg1 -A16M"   build-depends:       base     , futhark-    , text   default-language: Haskell2010  test-suite unit   type: exitcode-stdio-1.0   main-is: futhark_tests.hs   other-modules:+      Futhark.AD.DerivativesTests       Futhark.BenchTests       Futhark.Pkg.SolveTests       Futhark.IR.Prop.RearrangeTests

@@ -0,0 +1,23 @@+-- | Definitions related to automatic differentiation.+--+-- **Warning:** This is an experimental feature.  It has several known+-- bugs.  Documentation is very sparse.  You probably don't want to+-- use this unless you have the soul of a hero.++-- | Jacobian-Vector Product ("forward mode"), producing also the+-- primal result as the first element of the result tuple.+let jvp2 'a 'b (f: a -> b) (x: a) (x': a) : (b, b) =+  intrinsics.jvp2 (f, x, x')++-- | Vector-Jacobian Product ("reverse mode"), producing also the+-- primal result as the first element of the result tuple.+let vjp2 'a 'b (f: a -> b) (x: a) (y': b) : (b, a) =+  intrinsics.vjp2 (f, x, y')++-- | Jacobian-Vector Product ("forward mode").+let jvp 'a 'b (f: a -> b) (x: a) (x': a) : b =+  (jvp2 f x x').1++-- | Vector-Jacobian Product ("reverse mode").+let vjp 'a 'b (f: a -> b) (x: a) (y': b) : a =+  (vjp2 f x y').1

@@ -49,12 +49,12 @@ -- -- **Complexity:** O(1). def split [n] 't (i: i64) (xs: [n]t): ([i]t, []t) =-  (xs[:i] :> [i]t, xs[i:])+  (xs[0:i], xs[i:])  -- | Return the elements of the array in reverse order. -- -- **Complexity:** O(1).-def reverse [n] 't (x: [n]t): [n]t = x[::-1] :> [n]t+def reverse [n] 't (x: [n]t): [n]t = x[::-1]  -- | Concatenate two arrays.  Warning: never try to perform a reduction -- with this operator; it will not work.@@ -78,7 +78,7 @@ -- For example, if `b==rotate r a`, then `b[x] = a[x+r]`. -- -- **Complexity:** O(1).-def rotate [n] 't (r: i64) (xs: [n]t): [n]t = intrinsics.rotate (r, xs) :> *[n]t+def rotate [n] 't (r: i64) (xs: [n]t): [n]t = intrinsics.rotate (r, xs)  -- | Construct an array of consecutive integers of the given length, -- starting at 0.@@ -152,7 +152,7 @@ -- -- **Complexity:** O(1). def transpose [n] [m] 't (a: [n][m]t): [m][n]t =-  intrinsics.transpose a :> [m][n]t+  intrinsics.transpose a  -- | True if all of the input elements are true.  Produces true on an -- empty array.

@@ -5,6 +5,7 @@ open import "array" open import "math" open import "functional"+open import "ad"  -- | Create single-precision float from integer. def r32 (x: i32): f32 = f32.i32 x

@@ -48,7 +48,7 @@ -- -- **Span:** *O(S(f))* def map 'a [n] 'x (f: a -> x) (as: [n]a): *[n]x =-  intrinsics.map (f, as) :> *[n]x+  intrinsics.map (f, as)  -- | Apply the given function to each element of a single array. --@@ -150,7 +150,7 @@ -- -- **Span:** *O(log(n) ✕ W(op))* def scan [n] 'a (op: a -> a -> a) (ne: a) (as: [n]a): *[n]a =-  intrinsics.scan (op, ne, as) :> *[n]a+  intrinsics.scan (op, ne, as)  -- | Remove all those elements of `as` that do not satisfy the -- predicate `p`.@@ -219,7 +219,7 @@ -- -- **Span:** *O(S(f))* def map_stream [n] 'a 'b (f: (k: i64) -> [k]a -> [k]b) (as: [n]a): *[n]b =-  intrinsics.map_stream (f, as) :> *[n]b+  intrinsics.map_stream (f, as)  -- | Similar to `map_stream`@term, but the chunks do not necessarily -- correspond to subsequences of the original array (they may be@@ -229,7 +229,7 @@ -- -- **Span:** *O(S(f))* def map_stream_per [n] 'a 'b (f: (k: i64) -> [k]a -> [k]b) (as: [n]a): *[n]b =-  intrinsics.map_stream_per (f, as) :> *[n]b+  intrinsics.map_stream_per (f, as)  -- | Return `true` if the given function returns `true` for all -- elements in the array.@@ -278,7 +278,7 @@ -- -- **Span:** *O(1)* def scatter 't [m] [n] (dest: *[m]t) (is: [n]i64) (vs: [n]t): *[m]t =-  intrinsics.scatter (dest, is, vs) :> *[m]t+  intrinsics.scatter (dest, is, vs)  -- | `scatter_2d as is vs` is the equivalent of a `scatter` on a 2-dimensional -- array.@@ -287,7 +287,7 @@ -- -- **Span:** *O(1)* def scatter_2d 't [m] [n] [l] (dest: *[m][n]t) (is: [l](i64, i64)) (vs: [l]t): *[m][n]t =-  intrinsics.scatter_2d (dest, is, vs) :> *[m][n]t+  intrinsics.scatter_2d (dest, is, vs)  -- | `scatter_3d as is vs` is the equivalent of a `scatter` on a 3-dimensional -- array.@@ -296,4 +296,4 @@ -- -- **Span:** *O(1)* def scatter_3d 't [m] [n] [o] [l] (dest: *[m][n][o]t) (is: [l](i64, i64, i64)) (vs: [l]t): *[m][n][o]t =-  intrinsics.scatter_3d (dest, is, vs) :> *[m][n][o]t+  intrinsics.scatter_3d (dest, is, vs)

@@ -11,15 +11,15 @@ -- depended upon by soacs.fut.  So we just define a quick-and-dirty -- internal one here that uses the intrinsic version. local def internal_map 'a [n] 'x (f: a -> x) (as: [n]a): [n]x =-  intrinsics.map (f, as) :> [n]x+  intrinsics.map (f, as)  -- | Construct an array of pairs from two arrays. def zip [n] 'a 'b (as: [n]a) (bs: [n]b): *[n](a,b) =-  intrinsics.zip (as, bs) :> *[n](a,b)+  intrinsics.zip (as, bs)  -- | Construct an array of pairs from two arrays. def zip2 [n] 'a 'b (as: [n]a) (bs: [n]b): *[n](a,b) =-  zip as bs :> *[n](a,b)+  zip as bs  -- | As `zip2`@term, but with one more array. def zip3 [n] 'a 'b 'c (as: [n]a) (bs: [n]b) (cs: [n]c): *[n](a,b,c) =@@ -35,7 +35,7 @@  -- | Turn an array of pairs into two arrays. def unzip [n] 'a 'b (xs: [n](a,b)): ([n]a, [n]b) =-  intrinsics.unzip xs :> ([n]a, [n]b)+  intrinsics.unzip xs  -- | Turn an array of pairs into two arrays. def unzip2 [n] 'a 'b (xs: [n](a,b)): ([n]a, [n]b) =

@@ -1,176 +0,0 @@-// Start of chaselev.h--/* Implementation of Chase-lev's concurrent lock-free deque-   from ``Dynamic Circular Work-Stealing Deque`` (2005)-   This implementation was ported from-   https://github.com/deepsea-inria/heartbeat--   !!!-   This implementation leaks memory,-   if the circular array is grown-   as we don't maintain a list of the old buffers.-   However, we can't safely free it either as a stealing thread might-   be reading from it.-   !!!- */--#if defined(MCCHASELEV)-#include <stdlib.h>-#include <assert.h>-#include <string.h>---static struct subtask* const STEAL_RES_EMPTY = (struct subtask*) 0;-static struct subtask* const STEAL_RES_ABORT = (struct subtask*) 1;--static const int strong = 0;-static const int backoff_nb_cycles = 1l << 10;---static inline struct subtask* cb_get(struct subtask **buf, int64_t capacity, int64_t i)  {-  return (struct subtask*)__atomic_load_n(&buf[i % capacity], __ATOMIC_RELAXED);-}--static inline void cb_put (struct subtask **buf, int64_t capacity, int64_t i, struct subtask* x) {-  __atomic_store_n(&buf[i % capacity], x, __ATOMIC_RELAXED);-}--struct deque_buffer* grow(struct subtask **old_array,-                          int64_t old_capacity,-                          int64_t new_capacity,-                          int64_t b,-                          int64_t t)-{-  struct deque_buffer* new_deque_buffer = malloc(sizeof(struct deque_buffer));-  new_deque_buffer->size = new_capacity;-  new_deque_buffer->array = calloc(new_capacity,  sizeof(struct subtask*));--  for (int64_t i = t; i < b; i++) {-    cb_put(new_deque_buffer->array, new_capacity, i, cb_get(old_array, old_capacity, i));-  }-  return new_deque_buffer;-}--static inline int deque_init(struct deque *q, int64_t capacity) {-  assert(q != NULL);-  memset(q, 0, sizeof(struct deque));--  q->buffer = malloc(sizeof(struct deque_buffer));-  q->buffer->array = calloc(capacity, sizeof(struct subtask*));-  q->buffer->size = capacity;--  q->dead = 0;--  if (q->buffer->array == NULL) {-    return -1;-  }--  if (q->buffer == NULL) {-    return -1;-  }-  return 0;-}--static inline void deque_destroy(struct deque* q)-{-  q->dead = 1;-  free(q->buffer->array);-  free(q->buffer);-}--static inline int cas_top (struct deque *q, int64_t old_val, int64_t new_val) {-  int64_t ov = old_val;-  if(__atomic_compare_exchange_n(&q->top, &ov, new_val, strong,-                                 __ATOMIC_SEQ_CST, __ATOMIC_RELAXED)) {-    return 1;-  }-  spin_for(backoff_nb_cycles);-  return 0;-}---void push_back(struct deque *q, struct subtask*subtask)-{-  assert(subtask != NULL);-  assert(q != NULL);--  int64_t b = __atomic_load_n(&q->bottom, __ATOMIC_RELAXED); // load atomically-  int64_t t = __atomic_load_n(&q->top, __ATOMIC_ACQUIRE);    // load atomically-  struct deque_buffer *buffer = __atomic_load_n(&q->buffer, __ATOMIC_RELAXED);-  if (b-t >= (buffer->size - 1)) {-    // grow_queue-    struct subtask **old_array = buffer->array;-    int64_t old_capacity = __atomic_load_n(&buffer->size, __ATOMIC_RELAXED);-    int64_t new_capacity = old_capacity * 2;-    struct deque_buffer *new_buffer = grow(old_array, old_capacity, new_capacity, b, t);-    __atomic_store_n(&q->buffer, new_buffer, __ATOMIC_RELEASE);-    buffer = __atomic_load_n(&q->buffer, __ATOMIC_RELAXED);-    memset(old_array, 0, sizeof(struct subtask*) * old_capacity);-    // free(old_array);  Not safe!!-  }--  cb_put(buffer->array, buffer->size, b, subtask);-  __atomic_thread_fence(__ATOMIC_RELEASE);-  __atomic_store_n(&q->bottom, b+1, __ATOMIC_RELAXED);-  return;-}---struct subtask * pop_back(struct deque *q)-{-  int64_t b = __atomic_load_n(&q->bottom, __ATOMIC_RELAXED) - 1; // load atomically-  struct deque_buffer *buffer = __atomic_load_n(&q->buffer, __ATOMIC_RELAXED);-  __atomic_store_n(&q->bottom, b, __ATOMIC_RELAXED);-	__atomic_thread_fence(__ATOMIC_SEQ_CST);-  int64_t t = __atomic_load_n(&q->top, __ATOMIC_RELAXED);-  if (b < t) {-    __atomic_store_n(&q->bottom, t, __ATOMIC_RELAXED);-    return NULL;-  }-  struct subtask* item = cb_get(buffer->array, buffer->size, b);-  if (b > t) {-    return item;-  }--  // else there's only one item left-  // Did we win the race?-  if (!cas_top(q, t, t + 1)) {-    item = NULL;-  }-  __atomic_store_n(&q->bottom, t+1, __ATOMIC_RELAXED);-  return item;-}--struct subtask* steal(struct deque *q)-{-  assert(q != NULL);--  int64_t t = __atomic_load_n(&q->top, __ATOMIC_ACQUIRE);    // load atomically-  __atomic_thread_fence(__ATOMIC_SEQ_CST);-  int64_t b = __atomic_load_n(&q->bottom, __ATOMIC_ACQUIRE); // load atomically-  if (t >= b) {-    return STEAL_RES_EMPTY;-  }--  struct deque_buffer *buffer = __atomic_load_n(&q->buffer, __ATOMIC_CONSUME);-  struct subtask* item = cb_get(buffer->array, buffer->size, t);-  if (!cas_top(q, t, t + 1)) {-    return STEAL_RES_ABORT;-  }--  return item;-}---static inline size_t nb_subtasks(struct deque *q)-{-  return (size_t)__atomic_load_n(&q->bottom, __ATOMIC_RELAXED) - __atomic_load_n(&q->top, __ATOMIC_RELAXED);-}--static inline int empty(struct deque *q)-{-  return nb_subtasks(q) < 1;-}--#endif-// end of chaselev.h

@@ -290,7 +290,9 @@     { 7, 0, "compute_70" },     { 7, 2, "compute_72" },     { 7, 5, "compute_75" },-    { 8, 0, "compute_80" }+    { 8, 0, "compute_80" },+    { 8, 6, "compute_80" },+    { 8, 7, "compute_80" }   };    int major = device_query(dev, COMPUTE_CAPABILITY_MAJOR);

@@ -218,7 +218,12 @@          build_options += ["-DLOCKSTEP_WIDTH={}".format(lockstep_width)] -        build_options += ["-D{}={}".format(s.replace('z', 'zz').replace('.', 'zi').replace('#', 'zh'),v) for (s,v) in self.sizes.items()]+        build_options += ["-D{}={}".format(s.+                                           replace('z', 'zz').+                                           replace('.', 'zi').+                                           replace('#', 'zh').+                                           replace('\'', 'zq'),+                                           v) for (s,v) in self.sizes.items()]          if (self.platform.name == 'Oclgrind'):             build_options += ['-DEMULATE_F16']

@@ -0,0 +1,386 @@+{-# LANGUAGE OverloadedStrings #-}++-- | Partial derivatives of scalar Futhark operations and built-in functions.+module Futhark.AD.Derivatives+  ( pdBuiltin,+    pdBinOp,+    pdUnOp,+  )+where++import Data.Bifunctor (bimap)+import Futhark.Analysis.PrimExp.Convert+import Futhark.IR.Syntax.Core+import Futhark.Util.IntegralExp+import Prelude hiding (quot)++iConst :: IntType -> Integer -> PrimExp VName+iConst it x = ValueExp $ IntValue $ intValue it x++fConst :: FloatType -> Double -> PrimExp VName+fConst ft x = ValueExp $ FloatValue $ floatValue ft x++untyped2 :: (TPrimExp t v, TPrimExp t v) -> (PrimExp v, PrimExp v)+untyped2 = bimap untyped untyped++-- | @pdUnOp op x@ computes the partial derivatives of @op@+-- with respect to @x@.+pdUnOp :: UnOp -> PrimExp VName -> PrimExp VName+pdUnOp (Abs it) a = UnOpExp (SSignum it) a+pdUnOp (FAbs ft) a = UnOpExp (FSignum ft) a+pdUnOp Not x = x+pdUnOp (Complement it) x = UnOpExp (Complement it) x+pdUnOp (SSignum it) _ = iConst it 0+pdUnOp (USignum it) _ = iConst it 0+pdUnOp (FSignum ft) _ = fConst ft 0++type OnBinOp t v = TPrimExp t v -> TPrimExp t v -> (TPrimExp t v, TPrimExp t v)++intBinOp ::+  OnBinOp Int8 v ->+  OnBinOp Int16 v ->+  OnBinOp Int32 v ->+  OnBinOp Int64 v ->+  IntType ->+  PrimExp v ->+  PrimExp v ->+  (PrimExp v, PrimExp v)+intBinOp f _ _ _ Int8 a b = untyped2 $ f (isInt8 a) (isInt8 b)+intBinOp _ f _ _ Int16 a b = untyped2 $ f (isInt16 a) (isInt16 b)+intBinOp _ _ f _ Int32 a b = untyped2 $ f (isInt32 a) (isInt32 b)+intBinOp _ _ _ f Int64 a b = untyped2 $ f (isInt64 a) (isInt64 b)++floatBinOp ::+  OnBinOp Half v ->+  OnBinOp Float v ->+  OnBinOp Double v ->+  FloatType ->+  PrimExp v ->+  PrimExp v ->+  (PrimExp v, PrimExp v)+floatBinOp f _ _ Float16 a b = untyped2 $ f (isF16 a) (isF16 b)+floatBinOp _ f _ Float32 a b = untyped2 $ f (isF32 a) (isF32 b)+floatBinOp _ _ f Float64 a b = untyped2 $ f (isF64 a) (isF64 b)++-- | @pdBinOp op x y@ computes the partial derivatives of @op@ with+-- respect to @x@ and @y@.+pdBinOp :: BinOp -> PrimExp VName -> PrimExp VName -> (PrimExp VName, PrimExp VName)+pdBinOp (Add it _) _ _ = (iConst it 1, iConst it 1)+pdBinOp (Sub it _) _ _ = (iConst it 1, iConst it (-1))+pdBinOp (Mul _ _) x y = (y, x)+pdBinOp (Pow it) a b =+  intBinOp derivs derivs derivs derivs it a b+  where+    derivs x y = ((x `pow` (y - 1)) * y, 0) -- FIXME (wrt y)+pdBinOp (SDiv it _) a b =+  intBinOp derivs derivs derivs derivs it a b+  where+    derivs x y = (1 `quot` y, negate (x `quot` (y * y)))+pdBinOp (SDivUp it _) a b =+  intBinOp derivs derivs derivs derivs it a b+  where+    derivs x y = (1 `quot` y, negate (x `quot` (y * y)))+pdBinOp (SQuot it _) a b =+  intBinOp derivs derivs derivs derivs it a b+  where+    derivs x y = (1 `quot` y, negate (x `quot` (y * y)))+pdBinOp (UDiv it _) a b =+  intBinOp derivs derivs derivs derivs it a b+  where+    derivs x y = (1 `quot` y, negate (x `quot` (y * y)))+pdBinOp (UDivUp it _) a b =+  intBinOp derivs derivs derivs derivs it a b+  where+    derivs x y = (1 `quot` y, negate (x `quot` (y * y)))+pdBinOp (UMod it _) _ _ = (iConst it 1, iConst it 0) -- FIXME+pdBinOp (SMod it _) _ _ = (iConst it 1, iConst it 0) -- FIXME+pdBinOp (SRem it _) _ _ = (iConst it 1, iConst it 0) -- FIXME+pdBinOp (FMod ft) _ _ = (fConst ft 1, fConst ft 0) -- FIXME+pdBinOp (UMax it) a b =+  intBinOp derivs derivs derivs derivs it a b+  where+    derivs x y = (fromBoolExp (x .>=. y), fromBoolExp (x .<. y))+pdBinOp (SMax it) a b =+  intBinOp derivs derivs derivs derivs it a b+  where+    derivs x y = (fromBoolExp (x .>=. y), fromBoolExp (x .<. y))+pdBinOp (UMin it) a b =+  intBinOp derivs derivs derivs derivs it a b+  where+    derivs x y = (fromBoolExp (x .<=. y), fromBoolExp (x .>. y))+pdBinOp (SMin it) a b =+  intBinOp derivs derivs derivs derivs it a b+  where+    derivs x y = (fromBoolExp (x .<=. y), fromBoolExp (x .>. y))+--+pdBinOp (Shl it) a b =+  pdBinOp (Mul it OverflowWrap) a $ BinOpExp (Pow it) (iConst it 2) b+pdBinOp (LShr it) a b =+  pdBinOp (UDiv it Unsafe) a $ BinOpExp (Pow it) (iConst it 2) b+pdBinOp (AShr it) a b =+  pdBinOp (SDiv it Unsafe) a $ BinOpExp (Pow it) (iConst it 2) b+pdBinOp (And it) _a _b = (iConst it 0, iConst it 0) -- FIXME+pdBinOp (Or it) _a _b = (iConst it 0, iConst it 0) -- FIXME+pdBinOp (Xor it) _a _b = (iConst it 0, iConst it 0) -- FIXME+--+pdBinOp (FAdd ft) _ _ = (fConst ft 1, fConst ft 1)+pdBinOp (FSub ft) _ _ = (fConst ft 1, fConst ft (-1))+pdBinOp (FMul _) x y = (y, x)+pdBinOp (FDiv ft) a b =+  floatBinOp derivs derivs derivs ft a b+  where+    derivs x y = (1 / y, negate (x / (y * y)))+pdBinOp (FPow ft) a b =+  floatBinOp derivs derivs derivs ft a b+  where+    derivs x y = (y * (x ** (y - 1)), (x ** y) * log x)+pdBinOp (FMax ft) a b =+  floatBinOp derivs derivs derivs ft a b+  where+    derivs x y = (fromBoolExp (x .>=. y), fromBoolExp (x .<. y))+pdBinOp (FMin ft) a b =+  floatBinOp derivs derivs derivs ft a b+  where+    derivs x y = (fromBoolExp (x .<=. y), fromBoolExp (x .>. y))+pdBinOp LogAnd a b = (b, a)+pdBinOp LogOr _ _ = (ValueExp $ BoolValue True, ValueExp $ BoolValue False)++-- | @pdBuiltin f args i@ computes the partial derivative of @f@+-- applied to @args@ with respect to each of its arguments.  Returns+-- 'Nothing' if no such derivative is known.+pdBuiltin :: Name -> [PrimExp VName] -> Maybe [PrimExp VName]+pdBuiltin "sqrt16" [x] =+  Just [untyped $ 1 / (2 * sqrt (isF16 x))]+pdBuiltin "sqrt32" [x] =+  Just [untyped $ 1 / (2 * sqrt (isF32 x))]+pdBuiltin "sqrt64" [x] =+  Just [untyped $ 1 / (2 * sqrt (isF64 x))]+pdBuiltin "cbrt16" [x] =+  Just [untyped $ 1 / (3 * cbrt16 (isF16 x) * cbrt16 (isF16 x))]+  where+    cbrt16 a = isF16 $ FunExp "cbrt16" [untyped a] $ FloatType Float16+pdBuiltin "cbrt32" [x] =+  Just [untyped $ 1 / (3 * cbrt32 (isF32 x) * cbrt32 (isF32 x))]+  where+    cbrt32 a = isF32 $ FunExp "cbrt32" [untyped a] $ FloatType Float32+pdBuiltin "cbrt64" [x] =+  Just [untyped $ 1 / (3 * cbrt64 (isF64 x) * cbrt64 (isF64 x))]+  where+    cbrt64 a = isF64 $ FunExp "cbrt64" [untyped a] $ FloatType Float32+pdBuiltin "log16" [x] =+  Just [untyped $ 1 / isF16 x]+pdBuiltin "log32" [x] =+  Just [untyped $ 1 / isF32 x]+pdBuiltin "log64" [x] =+  Just [untyped $ 1 / isF64 x]+pdBuiltin "log10_16" [x] =+  Just [untyped $ 1 / (isF16 x * log 10)]+pdBuiltin "log10_32" [x] =+  Just [untyped $ 1 / (isF32 x * log 10)]+pdBuiltin "log10_64" [x] =+  Just [untyped $ 1 / (isF64 x * log 10)]+pdBuiltin "log2_16" [x] =+  Just [untyped $ 1 / (isF16 x * log 2)]+pdBuiltin "log2_32" [x] =+  Just [untyped $ 1 / (isF32 x * log 2)]+pdBuiltin "log2_64" [x] =+  Just [untyped $ 1 / (isF64 x * log 2)]+pdBuiltin "exp16" [x] =+  Just [untyped $ exp (isF16 x)]+pdBuiltin "exp32" [x] =+  Just [untyped $ exp (isF32 x)]+pdBuiltin "exp64" [x] =+  Just [untyped $ exp (isF64 x)]+pdBuiltin "sin16" [x] =+  Just [untyped $ cos (isF16 x)]+pdBuiltin "sin32" [x] =+  Just [untyped $ cos (isF32 x)]+pdBuiltin "sin64" [x] =+  Just [untyped $ cos (isF64 x)]+pdBuiltin "sinh16" [x] =+  Just [untyped $ cosh (isF16 x)]+pdBuiltin "sinh32" [x] =+  Just [untyped $ cosh (isF32 x)]+pdBuiltin "sinh64" [x] =+  Just [untyped $ cosh (isF64 x)]+pdBuiltin "cos16" [x] =+  Just [untyped $ -sin (isF16 x)]+pdBuiltin "cos32" [x] =+  Just [untyped $ -sin (isF32 x)]+pdBuiltin "cos64" [x] =+  Just [untyped $ -sin (isF64 x)]+pdBuiltin "cosh16" [x] =+  Just [untyped $ sinh (isF16 x)]+pdBuiltin "cosh32" [x] =+  Just [untyped $ sinh (isF32 x)]+pdBuiltin "cosh64" [x] =+  Just [untyped $ sinh (isF64 x)]+pdBuiltin "tan16" [x] =+  Just [untyped $ 1 / (cos (isF16 x) * cos (isF16 x))]+pdBuiltin "tan32" [x] =+  Just [untyped $ 1 / (cos (isF32 x) * cos (isF32 x))]+pdBuiltin "tan64" [x] =+  Just [untyped $ 1 / (cos (isF64 x) * cos (isF64 x))]+pdBuiltin "asin16" [x] =+  Just [untyped $ 1 / sqrt (1 - isF16 x * isF16 x)]+pdBuiltin "asin32" [x] =+  Just [untyped $ 1 / sqrt (1 - isF32 x * isF32 x)]+pdBuiltin "asin64" [x] =+  Just [untyped $ 1 / sqrt (1 - isF64 x * isF64 x)]+pdBuiltin "asinh16" [x] =+  Just [untyped $ 1 / sqrt (1 + isF16 x * isF16 x)]+pdBuiltin "asinh32" [x] =+  Just [untyped $ 1 / sqrt (1 + isF32 x * isF32 x)]+pdBuiltin "asinh64" [x] =+  Just [untyped $ 1 / sqrt (1 + isF64 x * isF64 x)]+pdBuiltin "acos16" [x] =+  Just [untyped $ -1 / sqrt (1 - isF16 x * isF16 x)]+pdBuiltin "acos32" [x] =+  Just [untyped $ -1 / sqrt (1 - isF32 x * isF32 x)]+pdBuiltin "acos64" [x] =+  Just [untyped $ -1 / sqrt (1 - isF64 x * isF64 x)]+pdBuiltin "acosh16" [x] =+  Just [untyped $ 1 / sqrt (isF16 x * isF16 x - 1)]+pdBuiltin "acosh32" [x] =+  Just [untyped $ 1 / sqrt (isF32 x * isF32 x - 1)]+pdBuiltin "acosh64" [x] =+  Just [untyped $ 1 / sqrt (isF64 x * isF64 x - 1)]+pdBuiltin "atan16" [x] =+  Just [untyped $ 1 / (1 + isF16 x * isF16 x)]+pdBuiltin "atan32" [x] =+  Just [untyped $ 1 / (1 + isF32 x * isF32 x)]+pdBuiltin "atan64" [x] =+  Just [untyped $ 1 / (1 + isF64 x * isF64 x)]+pdBuiltin "atanh16" [x] =+  Just [untyped $ cosh (isF16 x) * cosh (isF16 x)]+pdBuiltin "atanh32" [x] =+  Just [untyped $ cosh (isF32 x) * cosh (isF32 x)]+pdBuiltin "atanh64" [x] =+  Just [untyped $ cosh (isF64 x) * cosh (isF64 x)]+pdBuiltin "atan2_16" [x, y] =+  Just+    [ untyped $ -isF16 y / (isF16 x * isF16 x + isF16 y * isF16 y),+      untyped $ -isF16 x / (isF16 x * isF16 x + isF16 y * isF16 y)+    ]+pdBuiltin "atan2_32" [x, y] =+  Just+    [ untyped $ -isF32 y / (isF32 x * isF32 x + isF32 y * isF32 y),+      untyped $ -isF32 x / (isF32 x * isF32 x + isF32 y * isF32 y)+    ]+pdBuiltin "atan2_64" [x, y] =+  Just+    [ untyped $ -isF64 y / (isF64 x * isF64 x + isF64 y * isF64 y),+      untyped $ -isF64 x / (isF64 x * isF64 x + isF64 y * isF64 y)+    ]+pdBuiltin "tanh16" [x] =+  Just [untyped $ 1 - tanh (isF16 x) * tanh (isF16 x)]+pdBuiltin "tanh32" [x] =+  Just [untyped $ 1 - tanh (isF32 x) * tanh (isF32 x)]+pdBuiltin "tanh64" [x] =+  Just [untyped $ 1 - tanh (isF64 x) * tanh (isF64 x)]+pdBuiltin "fma16" [a, b, _c] =+  Just [b, a, fConst Float16 1]+pdBuiltin "fma32" [a, b, _c] =+  Just [b, a, fConst Float32 1]+pdBuiltin "fma64" [a, b, _c] =+  Just [b, a, fConst Float64 1]+pdBuiltin "mad16" [a, b, _c] =+  Just [b, a, fConst Float16 1]+pdBuiltin "mad32" [a, b, _c] =+  Just [b, a, fConst Float32 1]+pdBuiltin "mad64" [a, b, _c] =+  Just [b, a, fConst Float64 1]+pdBuiltin "from_bits16" [_] =+  Just [fConst Float16 1]+pdBuiltin "from_bits32" [_] =+  Just [fConst Float32 1]+pdBuiltin "from_bits64" [_] =+  Just [fConst Float64 1]+pdBuiltin "to_bits16" [_] =+  Just [iConst Int16 1]+pdBuiltin "to_bits32" [_] =+  Just [iConst Int32 1]+pdBuiltin "to_bits64" [_] =+  Just [iConst Int64 1]+pdBuiltin "hypot16" [x, y] =+  Just+    [ untyped $ isF16 x / isF16 (FunExp "hypot16" [x, y] $ FloatType Float16),+      untyped $ isF16 y / isF16 (FunExp "hypot16" [x, y] $ FloatType Float16)+    ]+pdBuiltin "hypot32" [x, y] =+  Just+    [ untyped $ isF32 x / isF32 (FunExp "hypot32" [x, y] $ FloatType Float32),+      untyped $ isF32 y / isF32 (FunExp "hypot32" [x, y] $ FloatType Float32)+    ]+pdBuiltin "hypot64" [x, y] =+  Just+    [ untyped $ isF64 x / isF64 (FunExp "hypot64" [x, y] $ FloatType Float64),+      untyped $ isF64 y / isF64 (FunExp "hypot64" [x, y] $ FloatType Float64)+    ]+pdBuiltin "lerp16" [v0, v1, t] =+  Just+    [ untyped $ 1 - fMax16 0 (fMin16 1 (isF16 t)),+      untyped $ fMax16 0 (fMin16 1 (isF16 t)),+      untyped $ isF16 v1 - isF16 v0+    ]+pdBuiltin "lerp32" [v0, v1, t] =+  Just+    [ untyped $ 1 - fMax32 0 (fMin32 1 (isF32 t)),+      untyped $ fMax32 0 (fMin32 1 (isF32 t)),+      untyped $ isF32 v1 - isF32 v0+    ]+pdBuiltin "lerp64" [v0, v1, t] =+  Just+    [ untyped $ 1 - fMax64 0 (fMin64 1 (isF64 t)),+      untyped $ fMax64 0 (fMin64 1 (isF64 t)),+      untyped $ isF64 v1 - isF64 v0+    ]+pdBuiltin "erf16" [z] =+  Just [untyped $ (2 / sqrt pi) * exp (negate (isF16 z * isF16 z))]+pdBuiltin "erf32" [z] =+  Just [untyped $ (2 / sqrt pi) * exp (negate (isF32 z * isF32 z))]+pdBuiltin "erf64" [z] =+  Just [untyped $ (2 / sqrt pi) * exp (negate (isF64 z * isF64 z))]+pdBuiltin "erfc16" [z] =+  Just [untyped $ negate $ (2 / sqrt pi) * exp (negate (isF16 z * isF16 z))]+pdBuiltin "erfc32" [z] =+  Just [untyped $ negate $ (2 / sqrt pi) * exp (negate (isF32 z * isF32 z))]+pdBuiltin "erfc64" [z] =+  Just [untyped $ negate $ (2 / sqrt pi) * exp (negate (isF64 z * isF64 z))]+-- More problematic derivatives follow below.+pdBuiltin "mul_hi8" [x, y] = Just [y, x]+pdBuiltin "mul_hi16" [x, y] = Just [y, x]+pdBuiltin "mul_hi32" [x, y] = Just [y, x]+pdBuiltin "mul_hi64" [x, y] = Just [y, x]+pdBuiltin "mad_hi8" [a, b, _c] = Just [b, a, iConst Int8 1]+pdBuiltin "mad_hi16" [a, b, _c] = Just [b, a, iConst Int16 1]+pdBuiltin "mad_hi32" [a, b, _c] = Just [b, a, iConst Int32 1]+pdBuiltin "mad_hi64" [a, b, _c] = Just [b, a, iConst Int64 1]+pdBuiltin "isnan16" [_] = Just [untyped false]+pdBuiltin "isnan32" [_] = Just [untyped false]+pdBuiltin "isnan64" [_] = Just [untyped false]+pdBuiltin "isinf16" [_] = Just [untyped false]+pdBuiltin "isinf32" [_] = Just [untyped false]+pdBuiltin "isinf64" [_] = Just [untyped false]+pdBuiltin "round16" [_] = Just [fConst Float16 0]+pdBuiltin "round32" [_] = Just [fConst Float32 0]+pdBuiltin "round64" [_] = Just [fConst Float64 0]+pdBuiltin "ceil16" [_] = Just [fConst Float16 0]+pdBuiltin "ceil32" [_] = Just [fConst Float32 0]+pdBuiltin "ceil64" [_] = Just [fConst Float64 0]+pdBuiltin "floor16" [_] = Just [fConst Float16 0]+pdBuiltin "floor32" [_] = Just [fConst Float32 0]+pdBuiltin "floor64" [_] = Just [fConst Float64 0]+pdBuiltin "clz8" [_] = Just [iConst Int32 0]+pdBuiltin "clz16" [_] = Just [iConst Int32 0]+pdBuiltin "clz32" [_] = Just [iConst Int32 0]+pdBuiltin "clz64" [_] = Just [iConst Int32 0]+pdBuiltin "ctz8" [_] = Just [iConst Int32 0]+pdBuiltin "ctz16" [_] = Just [iConst Int32 0]+pdBuiltin "ctz32" [_] = Just [iConst Int32 0]+pdBuiltin "ctz64" [_] = Just [iConst Int32 0]+pdBuiltin "popc8" [_] = Just [iConst Int32 0]+pdBuiltin "popc16" [_] = Just [iConst Int32 0]+pdBuiltin "popc32" [_] = Just [iConst Int32 0]+pdBuiltin "popc64" [_] = Just [iConst Int32 0]+pdBuiltin _ _ = Nothing

@@ -0,0 +1,486 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GeneralizedNewtypeDeriving #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UndecidableInstances #-}++module Futhark.AD.Fwd (fwdJVP) where++import Control.Monad+import Control.Monad.RWS.Strict+import Control.Monad.State.Strict+import Data.Bifunctor (second)+import qualified Data.Kind+import Data.List (transpose)+import Data.List.NonEmpty (NonEmpty (..))+import qualified Data.Map as M+import Futhark.AD.Derivatives+import Futhark.Analysis.PrimExp.Convert+import Futhark.Builder+import Futhark.Construct+import Futhark.IR.SOACS++zeroTan :: Type -> ADM SubExp+zeroTan (Prim t) = pure $ constant $ blankPrimValue t+zeroTan t = error $ "zeroTan on non-primitive type: " ++ pretty t++zeroExp :: Type -> Exp SOACS+zeroExp (Prim pt) =+  BasicOp $ SubExp $ Constant $ blankPrimValue pt+zeroExp (Array pt shape _) =+  BasicOp $ Replicate shape $ Constant $ blankPrimValue pt+zeroExp t = error $ "zeroExp: " ++ show t++tanType :: TypeBase s u -> ADM (TypeBase s u)+tanType (Acc acc ispace ts u) = do+  ts_tan <- mapM tanType ts+  pure $ Acc acc ispace (ts ++ ts_tan) u+tanType t = pure t++slocal' :: ADM a -> ADM a+slocal' = slocal id++slocal :: (RState -> RState) -> ADM a -> ADM a+slocal f m = do+  s <- get+  modify f+  a <- m+  modify $ \s' -> s' {stateTans = stateTans s}+  pure a++data RState = RState+  { stateTans :: M.Map VName VName,+    stateNameSource :: VNameSource+  }++newtype ADM a = ADM (BuilderT SOACS (State RState) a)+  deriving+    ( Functor,+      Applicative,+      Monad,+      MonadState RState,+      MonadFreshNames,+      HasScope SOACS,+      LocalScope SOACS+    )++instance MonadBuilder ADM where+  type Rep ADM = SOACS+  mkExpDecM pat e = ADM $ mkExpDecM pat e+  mkBodyM bnds res = ADM $ mkBodyM bnds res+  mkLetNamesM pat e = ADM $ mkLetNamesM pat e++  addStms = ADM . addStms+  collectStms (ADM m) = ADM $ collectStms m++instance MonadFreshNames (State RState) where+  getNameSource = gets stateNameSource+  putNameSource src = modify (\env -> env {stateNameSource = src})++runADM :: MonadFreshNames m => ADM a -> m a+runADM (ADM m) =+  modifyNameSource $ \vn ->+    second stateNameSource $+      runState+        (fst <$> runBuilderT m mempty)+        (RState mempty vn)++tanVName :: VName -> ADM VName+tanVName v = newVName (baseString v <> "_tan")++insertTan :: VName -> VName -> ADM ()+insertTan v v' =+  modify $ \env -> env {stateTans = M.insert v v' (stateTans env)}++class TanBuilder a where+  type Bundled a :: Data.Kind.Type+  type Bundled a = [a]+  newTan :: a -> ADM a+  bundleNew :: a -> ADM (Bundled a)++instance (Monoid (Bundled a), TanBuilder a) => TanBuilder [a] where+  type Bundled [a] = Bundled a+  newTan = mapM newTan+  bundleNew = fmap mconcat . mapM bundleNew++instance TanBuilder (PatElem (TypeBase s u)) where+  newTan (PatElem p t)+    | isAcc t = do+        insertTan p p+        t' <- tanType t+        pure $ PatElem p t'+    | otherwise = do+        p' <- tanVName p+        insertTan p p'+        t' <- tanType t+        pure $ PatElem p' t'+  bundleNew pe@(PatElem _ t) = do+    pe' <- newTan pe+    if isAcc t+      then pure [pe']+      else pure [pe, pe']++instance TanBuilder (Pat (TypeBase s u)) where+  type Bundled (Pat (TypeBase s u)) = Pat (TypeBase s u)+  newTan (Pat pes) = Pat <$> newTan pes+  bundleNew (Pat pes) = Pat <$> bundleNew pes++instance TanBuilder (Param (TypeBase s u)) where+  newTan (Param _ p t) = do+    PatElem p' t' <- newTan $ PatElem p t+    pure $ Param mempty p' t'+  bundleNew param@(Param _ _ (Prim Unit)) =+    pure [param]+  bundleNew param@(Param _ _ t) = do+    param' <- newTan param+    if isAcc t+      then pure [param']+      else pure [param, param']++instance Tangent a => TanBuilder (Param (TypeBase s u), a) where+  newTan (p, x) = (,) <$> newTan p <*> tangent x+  bundleNew (p, x) = do+    b <- bundleNew p+    x_tan <- tangent x+    pure $ zip b [x, x_tan]++class Tangent a where+  type BundledTan a :: Data.Kind.Type+  type BundledTan a = [a]+  tangent :: a -> ADM a+  bundleTan :: a -> ADM (BundledTan a)++instance Tangent (TypeBase s u) where+  tangent = tanType+  bundleTan t+    | isAcc t = do+        t' <- tangent t+        pure [t']+    | otherwise = do+        t' <- tangent t+        pure [t, t']++instance (Monoid (BundledTan a), Tangent a) => Tangent [a] where+  type BundledTan [a] = BundledTan a+  tangent = mapM tangent+  bundleTan = (mconcat <$>) . mapM bundleTan++instance Tangent VName where+  tangent v = do+    maybeTan <- gets $ M.lookup v . stateTans+    case maybeTan of+      Just v_tan -> pure v_tan+      Nothing -> do+        t <- lookupType v+        letExp (baseString v <> "_implicit_tan") $ zeroExp t+  bundleTan v = do+    t <- lookupType v+    if isAcc t+      then pure [v]+      else do+        v_tan <- tangent v+        pure [v, v_tan]++instance Tangent SubExp where+  tangent (Constant c) = zeroTan $ Prim $ primValueType c+  tangent (Var v) = Var <$> tangent v+  bundleTan c@Constant {} = do+    c_tan <- tangent c+    pure [c, c_tan]+  bundleTan (Var v) = fmap Var <$> bundleTan v++instance Tangent SubExpRes where+  tangent (SubExpRes cs se) = SubExpRes cs <$> tangent se+  bundleTan (SubExpRes cs se) = map (SubExpRes cs) <$> bundleTan se++basicFwd :: Pat Type -> StmAux () -> BasicOp -> ADM ()+basicFwd pat aux op = do+  pat_tan <- newTan pat+  case op of+    SubExp se -> do+      se_tan <- tangent se+      addStm $ Let pat_tan aux $ BasicOp $ SubExp se_tan+    Opaque opaqueop se -> do+      se_tan <- tangent se+      addStm $ Let pat_tan aux $ BasicOp $ Opaque opaqueop se_tan+    ArrayLit ses t -> do+      ses_tan <- tangent ses+      addStm $ Let pat_tan aux $ BasicOp $ ArrayLit ses_tan t+    UnOp unop x -> do+      let t = unOpType unop+          x_pe = primExpFromSubExp t x+          dx = pdUnOp unop x_pe+      x_tan <- primExpFromSubExp t <$> tangent x+      auxing aux $ letBindNames (patNames pat_tan) <=< toExp $ x_tan ~*~ dx+    BinOp bop x y -> do+      let t = binOpType bop+      x_tan <- primExpFromSubExp t <$> tangent x+      y_tan <- primExpFromSubExp t <$> tangent y+      let (wrt_x, wrt_y) =+            pdBinOp bop (primExpFromSubExp t x) (primExpFromSubExp t y)+      auxing aux $+        letBindNames (patNames pat_tan) <=< toExp $+          x_tan ~*~ wrt_x ~+~ y_tan ~*~ wrt_y+    CmpOp {} ->+      addStm $ Let pat_tan aux $ BasicOp op+    ConvOp cop x -> do+      x_tan <- tangent x+      addStm $ Let pat_tan aux $ BasicOp $ ConvOp cop x_tan+    Assert {} -> pure ()+    Index arr slice -> do+      arr_tan <- tangent arr+      addStm $ Let pat_tan aux $ BasicOp $ Index arr_tan slice+    Update safety arr slice se -> do+      arr_tan <- tangent arr+      se_tan <- tangent se+      addStm $ Let pat_tan aux $ BasicOp $ Update safety arr_tan slice se_tan+    Concat d (arr :| arrs) w -> do+      arr_tan <- tangent arr+      arrs_tans <- tangent arrs+      addStm $ Let pat_tan aux $ BasicOp $ Concat d (arr_tan :| arrs_tans) w+    Copy arr -> do+      arr_tan <- tangent arr+      addStm $ Let pat_tan aux $ BasicOp $ Copy arr_tan+    Manifest ds arr -> do+      arr_tan <- tangent arr+      addStm $ Let pat_tan aux $ BasicOp $ Manifest ds arr_tan+    Iota n _ _ it -> do+      addStm $ Let pat_tan aux $ BasicOp $ Replicate (Shape [n]) (intConst it 0)+    Replicate n x -> do+      x_tan <- tangent x+      addStm $ Let pat_tan aux $ BasicOp $ Replicate n x_tan+    Scratch t shape ->+      addStm $ Let pat_tan aux $ BasicOp $ Scratch t shape+    Reshape reshape arr -> do+      arr_tan <- tangent arr+      addStm $ Let pat_tan aux $ BasicOp $ Reshape reshape arr_tan+    Rearrange perm arr -> do+      arr_tan <- tangent arr+      addStm $ Let pat_tan aux $ BasicOp $ Rearrange perm arr_tan+    Rotate rots arr -> do+      arr_tan <- tangent arr+      addStm $ Let pat_tan aux $ BasicOp $ Rotate rots arr_tan+    _ -> error $ "basicFwd: Unsupported op " ++ pretty op++fwdLambda :: Lambda SOACS -> ADM (Lambda SOACS)+fwdLambda l@(Lambda params body ret) =+  Lambda <$> bundleNew params <*> inScopeOf l (fwdBody body) <*> bundleTan ret++fwdStreamLambda :: Lambda SOACS -> ADM (Lambda SOACS)+fwdStreamLambda l@(Lambda params body ret) =+  Lambda <$> ((take 1 params ++) <$> bundleNew (drop 1 params)) <*> inScopeOf l (fwdBody body) <*> bundleTan ret++interleave :: [a] -> [a] -> [a]+interleave xs ys = concat $ transpose [xs, ys]++zeroFromSubExp :: SubExp -> ADM VName+zeroFromSubExp (Constant c) =+  letExp "zero" $+    BasicOp $ SubExp $ Constant $ blankPrimValue $ primValueType c+zeroFromSubExp (Var v) = do+  t <- lookupType v+  letExp "zero" $ zeroExp t++fwdSOAC :: Pat Type -> StmAux () -> SOAC SOACS -> ADM ()+fwdSOAC pat aux (Screma size xs (ScremaForm scs reds f)) = do+  pat' <- bundleNew pat+  xs' <- bundleTan xs+  scs' <- mapM fwdScan scs+  reds' <- mapM fwdRed reds+  f' <- fwdLambda f+  addStm $ Let pat' aux $ Op $ Screma size xs' $ ScremaForm scs' reds' f'+  where+    fwdScan :: Scan SOACS -> ADM (Scan SOACS)+    fwdScan sc = do+      op' <- fwdLambda $ scanLambda sc+      neutral_tans <- mapM zeroFromSubExp $ scanNeutral sc+      pure $+        Scan+          { scanNeutral = scanNeutral sc `interleave` map Var neutral_tans,+            scanLambda = op'+          }+    fwdRed :: Reduce SOACS -> ADM (Reduce SOACS)+    fwdRed red = do+      op' <- fwdLambda $ redLambda red+      neutral_tans <- mapM zeroFromSubExp $ redNeutral red+      pure $+        Reduce+          { redComm = redComm red,+            redLambda = op',+            redNeutral = redNeutral red `interleave` map Var neutral_tans+          }+fwdSOAC pat aux (Stream size xs form nes lam) = do+  pat' <- bundleNew pat+  lam' <- fwdStreamLambda lam+  xs' <- bundleTan xs+  nes_tan <- mapM (fmap Var . zeroFromSubExp) nes+  let nes' = interleave nes nes_tan+  case form of+    Sequential ->+      addStm $ Let pat' aux $ Op $ Stream size xs' Sequential nes' lam'+    Parallel o comm lam0 -> do+      lam0' <- fwdLambda lam0+      let form' = Parallel o comm lam0'+      addStm $ Let pat' aux $ Op $ Stream size xs' form' nes' lam'+fwdSOAC pat aux (Hist w arrs ops bucket_fun) = do+  pat' <- bundleNew pat+  ops' <- mapM fwdHist ops+  bucket_fun' <- fwdHistBucket bucket_fun+  arrs' <- bundleTan arrs+  addStm $ Let pat' aux $ Op $ Hist w arrs' ops' bucket_fun'+  where+    n_indices = sum $ map (shapeRank . histShape) ops+    fwdBodyHist (Body _ stms res) = buildBody_ $ do+      mapM_ fwdStm stms+      let (res_is, res_vs) = splitAt n_indices res+      (res_is ++) <$> bundleTan res_vs+    fwdHistBucket l@(Lambda params body ret) =+      let (r_is, r_vs) = splitAt n_indices ret+       in Lambda+            <$> bundleNew params+            <*> inScopeOf l (fwdBodyHist body)+            <*> ((r_is ++) <$> bundleTan r_vs)++    fwdHist :: HistOp SOACS -> ADM (HistOp SOACS)+    fwdHist (HistOp shape rf dest nes op) = do+      dest' <- bundleTan dest+      nes_tan <- mapM (fmap Var . zeroFromSubExp) nes+      op' <- fwdLambda op+      pure $+        HistOp+          { histShape = shape,+            histRaceFactor = rf,+            histDest = dest',+            histNeutral = interleave nes nes_tan,+            histOp = op'+          }+fwdSOAC (Pat pes) aux (Scatter w ivs lam as) = do+  as_tan <- mapM (\(s, n, a) -> do a_tan <- tangent a; pure (s, n, a_tan)) as+  pes_tan <- newTan pes+  ivs' <- bundleTan ivs+  let (as_ws, as_ns, _as_vs) = unzip3 as+      n_indices = sum $ zipWith (*) as_ns $ map length as_ws+  lam' <- fwdScatterLambda n_indices lam+  let s = Let (Pat (pes ++ pes_tan)) aux $ Op $ Scatter w ivs' lam' $ as ++ as_tan+  addStm s+  where+    fwdScatterLambda :: Int -> Lambda SOACS -> ADM (Lambda SOACS)+    fwdScatterLambda n_indices (Lambda params body ret) = do+      params' <- bundleNew params+      ret_tan <- tangent $ drop n_indices ret+      body' <- fwdBodyScatter n_indices body+      let indices = concat $ replicate 2 $ take n_indices ret+          ret' = indices ++ drop n_indices ret ++ ret_tan+      pure $ Lambda params' body' ret'+    fwdBodyScatter :: Int -> Body SOACS -> ADM (Body SOACS)+    fwdBodyScatter n_indices (Body _ stms res) = do+      (res_tan, stms') <- collectStms $ do+        mapM_ fwdStm stms+        tangent $ drop n_indices res+      let indices = concat $ replicate 2 $ take n_indices res+          res' = indices ++ drop n_indices res ++ res_tan+      pure $ mkBody stms' res'+fwdSOAC _ _ JVP {} =+  error "fwdSOAC: nested JVP not allowed."+fwdSOAC _ _ VJP {} =+  error "fwdSOAC: nested VJP not allowed."++fwdStm :: Stm SOACS -> ADM ()+fwdStm (Let pat aux (BasicOp (UpdateAcc acc i x))) = do+  pat' <- bundleNew pat+  x' <- bundleTan x+  acc_tan <- tangent acc+  addStm $ Let pat' aux $ BasicOp $ UpdateAcc acc_tan i x'+fwdStm stm@(Let pat aux (BasicOp e)) = do+  -- XXX: this has to be too naive.+  unless (any isAcc $ patTypes pat) $+    addStm stm+  basicFwd pat aux e+fwdStm stm@(Let pat _ (Apply f args _ _))+  | Just (ret, argts) <- M.lookup f builtInFunctions = do+      addStm stm+      arg_tans <-+        zipWith primExpFromSubExp argts <$> mapM (tangent . fst) args+      pat_tan <- newTan pat+      let arg_pes = zipWith primExpFromSubExp argts (map fst args)+      case pdBuiltin f arg_pes of+        Nothing ->+          error $ "No partial derivative defined for builtin function: " ++ pretty f+        Just derivs -> do+          let convertTo tt e+                | e_t == tt = e+                | otherwise =+                    case (tt, e_t) of+                      (IntType tt', IntType ft) -> ConvOpExp (SExt ft tt') e+                      (FloatType tt', FloatType ft) -> ConvOpExp (FPConv ft tt') e+                      (Bool, FloatType ft) -> ConvOpExp (FToB ft) e+                      (FloatType tt', Bool) -> ConvOpExp (BToF tt') e+                      _ -> error $ "fwdStm.convertTo: " ++ pretty (f, tt, e_t)+                where+                  e_t = primExpType e+          zipWithM_ (letBindNames . pure) (patNames pat_tan)+            =<< mapM toExp (zipWith (~*~) (map (convertTo ret) arg_tans) derivs)+fwdStm (Let pat aux (If cond t f (IfDec ret ifsort))) = do+  t' <- slocal' $ fwdBody t+  f' <- slocal' $ fwdBody f+  pat' <- bundleNew pat+  ret' <- bundleTan ret+  addStm $ Let pat' aux $ If cond t' f' $ IfDec ret' ifsort+fwdStm (Let pat aux (DoLoop val_pats loop@(WhileLoop v) body)) = do+  val_pats' <- bundleNew val_pats+  pat' <- bundleNew pat+  body' <-+    localScope (scopeOfFParams (map fst val_pats) <> scopeOf loop) $+      slocal' $ fwdBody body+  addStm $ Let pat' aux $ DoLoop val_pats' (WhileLoop v) body'+fwdStm (Let pat aux (DoLoop val_pats loop@(ForLoop i it bound loop_vars) body)) = do+  pat' <- bundleNew pat+  val_pats' <- bundleNew val_pats+  loop_vars' <- bundleNew loop_vars+  body' <-+    localScope (scopeOfFParams (map fst val_pats) <> scopeOf loop) $+      slocal' $ fwdBody body+  addStm $ Let pat' aux $ DoLoop val_pats' (ForLoop i it bound loop_vars') body'+fwdStm (Let pat aux (WithAcc inputs lam)) = do+  inputs' <- forM inputs $ \(shape, arrs, op) -> do+    arrs_tan <- tangent arrs+    op' <- case op of+      Nothing -> pure Nothing+      Just (op_lam, nes) -> do+        nes_tan <- mapM (fmap Var . zeroFromSubExp) nes+        op_lam' <- fwdLambda op_lam+        case op_lam' of+          Lambda ps body ret -> do+            let op_lam'' = Lambda (removeIndexTans (shapeRank shape) ps) body ret+            pure $ Just (op_lam'', interleave nes nes_tan)+    pure (shape, arrs <> arrs_tan, op')+  pat' <- bundleNew pat+  lam' <- fwdLambda lam+  addStm $ Let pat' aux $ WithAcc inputs' lam'+  where+    removeIndexTans 0 ps = ps+    removeIndexTans i (p : _ : ps) = p : removeIndexTans (i - 1) ps+    removeIndexTans _ ps = ps+fwdStm (Let pat aux (Op soac)) = fwdSOAC pat aux soac+fwdStm stm =+  error $ "unhandled forward mode AD for Stm: " ++ pretty stm ++ "\n" ++ show stm++fwdBody :: Body SOACS -> ADM (Body SOACS)+fwdBody (Body _ stms res) = buildBody_ $ do+  mapM_ fwdStm stms+  bundleTan res++fwdBodyTansLast :: Body SOACS -> ADM (Body SOACS)+fwdBodyTansLast (Body _ stms res) = buildBody_ $ do+  mapM_ fwdStm stms+  (res <>) <$> tangent res++fwdJVP :: MonadFreshNames m => Scope SOACS -> Lambda SOACS -> m (Lambda SOACS)+fwdJVP scope l@(Lambda params body ret) =+  runADM . localScope scope . inScopeOf l $ do+    params_tan <- newTan params+    body_tan <- fwdBodyTansLast body+    ret_tan <- tangent ret+    pure $ Lambda (params ++ params_tan) body_tan (ret <> ret_tan)

@@ -0,0 +1,444 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeFamilies #-}++-- Naming scheme:+--+-- An adjoint-related object for "x" is named "x_adj".  This means+-- both actual adjoints and statements.+--+-- Do not assume "x'" means anything related to derivatives.+module Futhark.AD.Rev (revVJP) where++import Control.Monad+import Data.List ((\\))+import Data.List.NonEmpty (NonEmpty (..))+import qualified Data.Map as M+import Futhark.AD.Derivatives+import Futhark.AD.Rev.Loop+import Futhark.AD.Rev.Monad+import Futhark.AD.Rev.SOAC+import Futhark.Analysis.PrimExp.Convert+import Futhark.Builder+import Futhark.IR.SOACS+import Futhark.Tools+import Futhark.Transform.Rename+import Futhark.Transform.Substitute+import Futhark.Util (takeLast)++patName :: Pat Type -> ADM VName+patName (Pat [pe]) = pure $ patElemName pe+patName pat = error $ "Expected single-element pattern: " ++ pretty pat++-- The vast majority of BasicOps require no special treatment in the+-- forward pass and produce one value (and hence one adjoint).  We+-- deal with that case here.+commonBasicOp :: Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM (VName, VName)+commonBasicOp pat aux op m = do+  addStm $ Let pat aux $ BasicOp op+  m+  pat_v <- patName pat+  pat_adj <- lookupAdjVal pat_v+  pure (pat_v, pat_adj)++diffBasicOp :: Pat Type -> StmAux () -> BasicOp -> ADM () -> ADM ()+diffBasicOp pat aux e m =+  case e of+    CmpOp cmp x y -> do+      (_pat_v, pat_adj) <- commonBasicOp pat aux e m+      returnSweepCode $ do+        let t = cmpOpType cmp+            update contrib = do+              void $ updateSubExpAdj x contrib+              void $ updateSubExpAdj y contrib++        case t of+          FloatType ft ->+            update <=< letExp "contrib" $+              If+                (Var pat_adj)+                (resultBody [constant (floatValue ft (1 :: Int))])+                (resultBody [constant (floatValue ft (0 :: Int))])+                (IfDec [Prim (FloatType ft)] IfNormal)+          IntType it ->+            update <=< letExp "contrib" $ BasicOp $ ConvOp (BToI it) (Var pat_adj)+          Bool ->+            update pat_adj+          Unit ->+            pure ()+    --+    ConvOp op x -> do+      (_pat_v, pat_adj) <- commonBasicOp pat aux e m+      returnSweepCode $ do+        contrib <-+          letExp "contrib" $ BasicOp $ ConvOp (flipConvOp op) $ Var pat_adj+        updateSubExpAdj x contrib+    --+    UnOp op x -> do+      (_pat_v, pat_adj) <- commonBasicOp pat aux e m++      returnSweepCode $ do+        let t = unOpType op+        contrib <- do+          let x_pe = primExpFromSubExp t x+              pat_adj' = primExpFromSubExp t (Var pat_adj)+              dx = pdUnOp op x_pe+          letExp "contrib" <=< toExp $ pat_adj' ~*~ dx++        updateSubExpAdj x contrib+    --+    BinOp op x y -> do+      (_pat_v, pat_adj) <- commonBasicOp pat aux e m++      returnSweepCode $ do+        let t = binOpType op+            (wrt_x, wrt_y) =+              pdBinOp op (primExpFromSubExp t x) (primExpFromSubExp t y)++            pat_adj' = primExpFromSubExp t $ Var pat_adj++        adj_x <- letExp "binop_x_adj" <=< toExp $ pat_adj' ~*~ wrt_x+        adj_y <- letExp "binop_y_adj" <=< toExp $ pat_adj' ~*~ wrt_y+        updateSubExpAdj x adj_x+        updateSubExpAdj y adj_y+    --+    SubExp se -> do+      (_pat_v, pat_adj) <- commonBasicOp pat aux e m+      returnSweepCode $ updateSubExpAdj se pat_adj+    --+    Assert {} ->+      void $ commonBasicOp pat aux e m+    --+    ArrayLit elems _ -> do+      (_pat_v, pat_adj) <- commonBasicOp pat aux e m+      t <- lookupType pat_adj+      returnSweepCode $ do+        forM_ (zip [(0 :: Int64) ..] elems) $ \(i, se) -> do+          let slice = fullSlice t [DimFix (constant i)]+          updateSubExpAdj se <=< letExp "elem_adj" $ BasicOp $ Index pat_adj slice+    --+    Index arr slice -> do+      (_pat_v, pat_adj) <- commonBasicOp pat aux e m+      returnSweepCode $ do+        void $ updateAdjSlice slice arr pat_adj+    FlatIndex {} -> error "FlatIndex not handled by AD yet."+    FlatUpdate {} -> error "FlatUpdate not handled by AD yet."+    --+    Opaque _ se -> do+      (_pat_v, pat_adj) <- commonBasicOp pat aux e m+      returnSweepCode $ updateSubExpAdj se pat_adj+    --+    Reshape _ arr -> do+      (_pat_v, pat_adj) <- commonBasicOp pat aux e m+      returnSweepCode $ do+        arr_dims <- arrayDims <$> lookupType arr+        void $+          updateAdj arr <=< letExp "adj_reshape" $+            BasicOp $ Reshape (map DimNew arr_dims) pat_adj+    --+    Rearrange perm arr -> do+      (_pat_v, pat_adj) <- commonBasicOp pat aux e m+      returnSweepCode $+        void $+          updateAdj arr <=< letExp "adj_rearrange" $+            BasicOp $ Rearrange (rearrangeInverse perm) pat_adj+    --+    Rotate rots arr -> do+      (_pat_v, pat_adj) <- commonBasicOp pat aux e m+      returnSweepCode $ do+        let neg = BasicOp . BinOp (Sub Int64 OverflowWrap) (intConst Int64 0)+        rots' <- mapM (letSubExp "rot_neg" . neg) rots+        void $+          updateAdj arr <=< letExp "adj_rotate" $+            BasicOp $ Rotate rots' pat_adj+    --+    Replicate (Shape ns) x -> do+      (_pat_v, pat_adj) <- commonBasicOp pat aux e m+      returnSweepCode $ do+        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 (baseString pat_adj <> "_flat") $+            BasicOp $ Reshape (map DimNew $ n : arrayDims x_t) pat_adj+        reduce <- reduceSOAC [Reduce Commutative lam [ne]]+        updateSubExpAdj x+          =<< letExp "rep_contrib" (Op $ Screma n [pat_adj_flat] reduce)+    --+    Concat d (arr :| arrs) _ -> do+      (_pat_v, pat_adj) <- commonBasicOp pat aux e m+      returnSweepCode $ do+        let sliceAdj _ [] = pure []+            sliceAdj start (v : vs) = do+              v_t <- lookupType v+              let w = arraySize 0 v_t+                  slice = DimSlice start w (intConst Int64 1)+              pat_adj_slice <-+                letExp (baseString pat_adj <> "_slice") $+                  BasicOp $ Index pat_adj (sliceAt v_t d [slice])+              start' <- letSubExp "start" $ BasicOp $ BinOp (Add Int64 OverflowUndef) start w+              slices <- sliceAdj start' vs+              pure $ pat_adj_slice : slices++        slices <- sliceAdj (intConst Int64 0) $ arr : arrs++        zipWithM_ updateAdj (arr : arrs) slices+    --+    Copy se -> do+      (_pat_v, pat_adj) <- commonBasicOp pat aux e m+      returnSweepCode $ void $ updateAdj se pat_adj+    --+    Manifest _ se -> do+      (_pat_v, pat_adj) <- commonBasicOp pat aux e m+      returnSweepCode $ void $ updateAdj se pat_adj+    --+    Scratch {} ->+      void $ commonBasicOp pat aux e m+    --+    Iota n _ _ t -> do+      (_pat_v, pat_adj) <- commonBasicOp pat aux e m+      returnSweepCode $ do+        ne <- letSubExp "zero" $ zeroExp $ Prim $ IntType t+        lam <- addLambda $ Prim $ IntType t+        reduce <- reduceSOAC [Reduce Commutative lam [ne]]+        updateSubExpAdj n+          =<< letExp "iota_contrib" (Op $ Screma n [pat_adj] reduce)+    --+    Update safety arr slice v -> do+      (_pat_v, pat_adj) <- commonBasicOp pat aux e m+      returnSweepCode $ do+        v_adj <- letExp "update_val_adj" $ BasicOp $ Index pat_adj slice+        t <- lookupType v_adj+        v_adj_copy <-+          case t of+            Array {} -> letExp "update_val_adj_copy" $ BasicOp $ Copy v_adj+            _ -> pure v_adj+        updateSubExpAdj v v_adj_copy+        zeroes <- letSubExp "update_zero" . zeroExp =<< subExpType v+        void $+          updateAdj arr+            =<< letExp "update_src_adj" (BasicOp $ Update safety pat_adj slice zeroes)+    -- See Note [Adjoints of accumulators]+    UpdateAcc _ is vs -> do+      addStm $ Let pat aux $ BasicOp e+      m+      pat_adjs <- mapM lookupAdjVal (patNames pat)+      returnSweepCode $ do+        forM_ (zip pat_adjs vs) $ \(adj, v) -> do+          adj_i <- letExp "updateacc_val_adj" $ BasicOp $ Index adj $ Slice $ map DimFix is+          updateSubExpAdj v adj_i++vjpOps :: VjpOps+vjpOps = VjpOps diffLambda diffStm++diffStm :: Stm SOACS -> ADM () -> ADM ()+diffStm (Let pat aux (BasicOp e)) m =+  diffBasicOp pat aux e m+diffStm stm@(Let pat _ (Apply f args _ _)) m+  | Just (ret, argts) <- M.lookup f builtInFunctions = do+      addStm stm+      m++      pat_adj <- lookupAdjVal =<< patName pat+      let arg_pes = zipWith primExpFromSubExp argts (map fst args)+          pat_adj' = primExpFromSubExp ret (Var pat_adj)+          convert ft tt+            | ft == tt = id+          convert (IntType ft) (IntType tt) = ConvOpExp (SExt ft tt)+          convert (FloatType ft) (FloatType tt) = ConvOpExp (FPConv ft tt)+          convert Bool (FloatType tt) = ConvOpExp (BToF tt)+          convert (FloatType ft) Bool = ConvOpExp (FToB ft)+          convert ft tt = error $ "diffStm.convert: " ++ pretty (f, ft, tt)++      contribs <-+        case pdBuiltin f arg_pes of+          Nothing ->+            error $ "No partial derivative defined for builtin function: " ++ pretty f+          Just derivs ->+            forM (zip derivs argts) $ \(deriv, argt) ->+              letExp "contrib" <=< toExp . convert ret argt $ pat_adj' ~*~ deriv++      zipWithM_ updateSubExpAdj (map fst args) contribs+diffStm stm@(Let pat _ (If cond tbody fbody _)) m = do+  addStm stm+  m+  returnSweepCode $ do+    let tbody_free = freeIn tbody+        fbody_free = freeIn fbody+        branches_free = namesToList $ tbody_free <> fbody_free++    adjs <- mapM lookupAdj $ patNames pat++    branches_free_adj <-+      ( pure . takeLast (length branches_free)+          <=< letTupExp "branch_adj"+          <=< renameExp+        )+        =<< eIf+          (eSubExp cond)+          (diffBody adjs branches_free tbody)+          (diffBody adjs branches_free fbody)+    zipWithM_ insAdj branches_free branches_free_adj+diffStm (Let pat aux (Op soac)) m =+  vjpSOAC vjpOps pat aux soac m+diffStm (Let pat aux loop@DoLoop {}) m =+  diffLoop diffStms pat aux loop m+-- See Note [Adjoints of accumulators]+diffStm stm@(Let pat _aux (WithAcc inputs lam)) m = do+  addStm stm+  m+  returnSweepCode $ do+    adjs <- mapM lookupAdj $ patNames pat+    lam' <- renameLambda lam+    free_vars <- filterM isActive $ namesToList $ freeIn lam'+    free_accs <- filterM (fmap isAcc . lookupType) free_vars+    let free_vars' = free_vars \\ free_accs+    lam'' <- diffLambda' adjs free_vars' lam'+    inputs' <- mapM renameInputLambda inputs+    free_adjs <- letTupExp "with_acc_contrib" $ WithAcc inputs' lam''+    zipWithM_ insAdj (arrs <> free_vars') free_adjs+  where+    arrs = concatMap (\(_, as, _) -> as) inputs+    renameInputLambda (shape, as, Just (f, nes)) = do+      f' <- renameLambda f+      pure (shape, as, Just (f', nes))+    renameInputLambda input = pure input+    diffLambda' res_adjs get_adjs_for (Lambda params body ts) =+      localScope (scopeOfLParams params) $ do+        Body () stms res <- diffBody res_adjs get_adjs_for body+        let body' = Body () stms $ take (length inputs) res <> takeLast (length get_adjs_for) res+        ts' <- mapM lookupType get_adjs_for+        pure $ Lambda params body' $ take (length inputs) ts <> ts'+diffStm stm _ = error $ "diffStm unhandled:\n" ++ pretty stm++diffStms :: Stms SOACS -> ADM ()+diffStms all_stms+  | Just (stm, stms) <- stmsHead all_stms = do+      (subst, copy_stms) <- copyConsumedArrsInStm stm+      let (stm', stms') = substituteNames subst (stm, stms)+      diffStms copy_stms >> diffStm stm' (diffStms stms')+      forM_ (M.toList subst) $ \(from, to) ->+        setAdj from =<< lookupAdj to+  | otherwise =+      pure ()++-- | Preprocess statements before differentiating.+-- For now, it's just stripmining.+preprocess :: Stms SOACS -> ADM (Stms SOACS)+preprocess = stripmineStms++diffBody :: [Adj] -> [VName] -> Body SOACS -> ADM (Body SOACS)+diffBody res_adjs get_adjs_for (Body () stms res) = subAD $+  subSubsts $ do+    let onResult (SubExpRes _ (Constant _)) _ = pure ()+        onResult (SubExpRes _ (Var v)) v_adj = void $ updateAdj v =<< adjVal v_adj+    (adjs, stms') <- collectStms $ do+      zipWithM_ onResult (takeLast (length res_adjs) res) res_adjs+      diffStms =<< preprocess stms+      mapM lookupAdjVal get_adjs_for+    pure $ Body () stms' $ res <> varsRes adjs++diffLambda :: [Adj] -> [VName] -> Lambda SOACS -> ADM (Lambda SOACS)+diffLambda res_adjs get_adjs_for (Lambda params body _) =+  localScope (scopeOfLParams params) $ do+    Body () stms res <- diffBody res_adjs get_adjs_for body+    let body' = Body () stms $ takeLast (length get_adjs_for) res+    ts' <- mapM lookupType get_adjs_for+    pure $ Lambda params body' ts'++revVJP :: MonadFreshNames m => Scope SOACS -> Lambda SOACS -> m (Lambda SOACS)+revVJP scope (Lambda params body ts) =+  runADM . localScope (scope <> scopeOfLParams params) $ do+    params_adj <- forM (zip (map resSubExp (bodyResult body)) ts) $ \(se, t) ->+      Param mempty <$> maybe (newVName "const_adj") adjVName (subExpVar se) <*> pure t++    body' <-+      localScope (scopeOfLParams params_adj) $+        diffBody+          (map adjFromParam params_adj)+          (map paramName params)+          body++    pure $ Lambda (params ++ params_adj) body' (ts <> map paramType params)++-- Note [Adjoints of accumulators]+--+-- 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+-- the result of previous invocations of VJP.  This means we can rely+-- on the operator having a constant adjoint (it's some kind of+-- addition).+--+-- 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+-- adjoint of 'acc(c, [n], {f64,f32})' should be two arrays of type+-- '[]f64', '[]f32'.  Our current design assumes that adjoints are+-- single variables.  This is fixable.+--+-- # Adjoint of UpdateAcc+--+--   Consider primal code+--+--     update_acc(acc, i, v)+--+--   Interpreted as an imperative statement, this means+--+--     acc[i] ⊕= v+--+--   for some '⊕'.  Normally all the compiler knows of '⊕' is that it+--   is associative and commutative, but because we assume that all+--   accumulators are the result of previous AD transformations, we+--   can assume that '⊕' actually behaves like addition - that is, has+--   unit partial derivatives.  So the return sweep is+--+--     v += acc_adj[i]+--+-- # Adjoint of Map+--+-- Suppose we have primal code+--+--   let acc' =+--     map (...) acc+--+-- where "acc : acc(c, [n], {f64})" and the width of the Map is "w".+-- Our normal transformation for Map input arrays is to similarly map+-- their adjoint, but clearly this doesn't work here because the+-- semantics of mapping an adjoint is an "implicit replicate".  So+-- when generating the return sweep we actually perform that+-- replication:+--+--   map (...) (replicate w acc_adj)+--+-- But what about the contributions to "acc'"?  Those we also have to+-- take special care of.  The result of the map itself is actually a+-- multidimensional array:+--+--   let acc_contribs =+--     map (...) (replicate w acc'_adj)+--+-- which we must then sum to add to the contribution.+--+--   acc_adj += sum(acc_contribs)+--+-- I'm slightly worried about the asymptotics of this, since my+-- intuition of this is that the contributions might be rather sparse.+-- (Maybe completely zero?  If so it will be simplified away+-- entirely.)  Perhaps a better solution is to treat+-- accumulator-inputs in the primal code as we do free variables, and+-- create accumulators for them in the return sweep.+--+-- # Consumption+--+-- A minor problem is that our usual way of handling consumption (Note+-- [Consumption]) is not viable, because accumulators are not+-- copyable.  Fortunately, while the accumulators that are consumed in+-- the forward sweep will also be present in the return sweep given+-- 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.

@@ -0,0 +1,456 @@+{-# LANGUAGE DeriveTraversable #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeFamilies #-}++module Futhark.AD.Rev.Loop (diffLoop, stripmineStms) where++import Control.Monad+import Data.Foldable (toList)+import Data.List (nub, (\\))+import qualified Data.Map as M+import Data.Maybe+import Futhark.AD.Rev.Monad+import qualified Futhark.Analysis.Alias as Alias+import Futhark.Analysis.PrimExp.Convert+import Futhark.Builder+import Futhark.IR.Aliases (consumedInStms)+import Futhark.IR.SOACS+import Futhark.Tools+import Futhark.Transform.Rename+import Futhark.Transform.Substitute+import Futhark.Util (traverseFold)++-- | A convenience function to bring the components of a for-loop into+-- scope and throw an error if the passed 'Exp' is not a for-loop.+bindForLoop ::+  PrettyRep rep =>+  Exp rep ->+  ( [(Param (FParamInfo rep), SubExp)] ->+    LoopForm rep ->+    VName ->+    IntType ->+    SubExp ->+    [(Param (LParamInfo rep), VName)] ->+    Body rep ->+    a+  ) ->+  a+bindForLoop (DoLoop val_pats form@(ForLoop i it bound loop_vars) body) f =+  f val_pats form i it bound loop_vars body+bindForLoop e _ = error $ "bindForLoop: not a for-loop:\n" <> pretty e++-- | A convenience function to rename a for-loop and then bind the+-- renamed components.+renameForLoop ::+  (MonadFreshNames m, Renameable rep, PrettyRep rep) =>+  Exp rep ->+  ( Exp rep ->+    [(Param (FParamInfo rep), SubExp)] ->+    LoopForm rep ->+    VName ->+    IntType ->+    SubExp ->+    [(Param (LParamInfo rep), VName)] ->+    Body rep ->+    m a+  ) ->+  m a+renameForLoop loop f = renameExp loop >>= \loop' -> bindForLoop loop' (f loop')++-- | Is the loop a while-loop?+isWhileLoop :: Exp rep -> Bool+isWhileLoop (DoLoop _ WhileLoop {} _) = True+isWhileLoop _ = False++-- | Transforms a 'ForLoop' into a 'ForLoop' with an empty list of+-- loop variables.+removeLoopVars :: MonadBuilder m => Exp (Rep m) -> m (Exp (Rep m))+removeLoopVars loop =+  bindForLoop loop $ \val_pats form i _it _bound loop_vars body -> do+    let indexify (x_param, xs) = do+          xs_t <- lookupType xs+          x' <-+            letExp (baseString $ paramName x_param) $+              BasicOp $ Index xs $ fullSlice xs_t [DimFix (Var i)]+          pure (paramName x_param, x')+    (substs_list, subst_stms) <- collectStms $ mapM indexify loop_vars+    let Body aux' stms' res' = substituteNames (M.fromList substs_list) body+    pure $ DoLoop val_pats form $ Body aux' (subst_stms <> stms') res'++-- | Augments a while-loop to also compute the number of iterations.+computeWhileIters :: Exp SOACS -> ADM SubExp+computeWhileIters (DoLoop val_pats (WhileLoop b) body) = do+  bound_v <- newVName "bound"+  let t = Prim $ IntType Int64+      bound_param = Param mempty bound_v t+  bound_init <- letSubExp "bound_init" $ zeroExp t+  body' <- localScope (scopeOfFParams [bound_param]) $+    buildBody_ $ do+      bound_plus_one <-+        let one = Constant $ IntValue $ intValue Int64 (1 :: Int)+         in letSubExp "bound+1" $ BasicOp $ BinOp (Add Int64 OverflowUndef) (Var bound_v) one+      addStms $ bodyStms body+      pure (pure (subExpRes bound_plus_one) <> bodyResult body)+  res <- letTupExp' "loop" $ DoLoop ((bound_param, bound_init) : val_pats) (WhileLoop b) body'+  pure $ head res+computeWhileIters e = error $ "convertWhileIters: not a while-loop:\n" <> pretty e++-- | Converts a 'WhileLoop' into a 'ForLoop'. Requires that the+-- surrounding 'DoLoop' is annotated with a @#[bound(n)]@ attribute,+-- where @n@ is an upper bound on the number of iterations of the+-- while-loop. The resulting for-loop will execute for @n@ iterations on+-- all inputs, so the tighter the bound the better.+convertWhileLoop :: SubExp -> Exp SOACS -> ADM (Exp SOACS)+convertWhileLoop bound_se (DoLoop val_pats (WhileLoop cond) body) =+  localScope (scopeOfFParams $ map fst val_pats) $ do+    i <- newVName "i"+    body' <-+      eBody+        [ eIf+            (pure $ BasicOp $ SubExp $ Var cond)+            (pure body)+            (resultBodyM $ map (Var . paramName . fst) val_pats)+        ]+    pure $ DoLoop val_pats (ForLoop i Int64 bound_se mempty) body'+convertWhileLoop _ e = error $ "convertWhileLoopBound: not a while-loop:\n" <> pretty e++-- | @nestifyLoop n bound loop@ transforms a loop into a depth-@n@ loop nest+-- of @bound@-iteration loops. This transformation does not preserve+-- the original semantics of the loop: @n@ and @bound@ may be arbitrary and have+-- no relation to the number of iterations of @loop@.+nestifyLoop ::+  SubExp ->+  Integer ->+  Exp SOACS ->+  ADM (Exp SOACS)+nestifyLoop bound_se = nestifyLoop' bound_se+  where+    nestifyLoop' offset n loop = bindForLoop loop nestify+      where+        nestify val_pats _form i it _bound loop_vars body+          | n > 1 = do+              renameForLoop loop $ \_loop' val_pats' _form' i' it' _bound' loop_vars' body' -> do+                let loop_params = map fst val_pats+                    loop_params' = map fst val_pats'+                    loop_inits' = map (Var . paramName) loop_params+                    val_pats'' = zip loop_params' loop_inits'+                outer_body <-+                  buildBody_ $+                    do+                      offset' <-+                        letSubExp "offset" $+                          BasicOp $ BinOp (Mul it OverflowUndef) offset (Var i)++                      inner_body <- insertStmsM $ do+                        i_inner <-+                          letExp "i_inner" $+                            BasicOp $ BinOp (Add it OverflowUndef) offset' (Var i')+                        pure $ substituteNames (M.singleton i' i_inner) body'++                      inner_loop <-+                        letTupExp "inner_loop"+                          =<< nestifyLoop'+                            offset'+                            (n - 1)+                            (DoLoop val_pats'' (ForLoop i' it' bound_se loop_vars') inner_body)+                      pure $ varsRes inner_loop+                pure $+                  DoLoop val_pats (ForLoop i it bound_se loop_vars) outer_body+          | n == 1 =+              pure $ DoLoop val_pats (ForLoop i it bound_se loop_vars) body+          | otherwise = pure loop++-- | @stripmine n pat loop@ stripmines a loop into a depth-@n@ loop nest.+-- An additional @bound - (floor(bound^(1/n)))^n@-iteration remainder loop is+-- inserted after the stripmined loop which executes the remaining iterations+-- so that the stripmined loop is semantically equivalent to the original loop.+stripmine :: Integer -> Pat Type -> Exp SOACS -> ADM (Stms SOACS)+stripmine n pat loop = do+  loop' <- removeLoopVars loop+  bindForLoop loop' $ \_val_pats _form _i it bound _loop_vars _body -> do+    let n_root = Constant $ FloatValue $ floatValue Float64 (1 / fromIntegral n :: Double)+    bound_float <- letSubExp "bound_f64" $ BasicOp $ ConvOp (UIToFP it Float64) bound+    bound' <- letSubExp "bound" $ BasicOp $ BinOp (FPow Float64) bound_float n_root+    bound_int <- letSubExp "bound_int" $ BasicOp $ ConvOp (FPToUI Float64 it) bound'+    total_iters <-+      letSubExp "total_iters" $+        BasicOp $ BinOp (Pow it) bound_int (Constant $ IntValue $ intValue it n)+    remain_iters <-+      letSubExp "remain_iters" $ BasicOp $ BinOp (Sub it OverflowUndef) bound total_iters+    mined_loop <- nestifyLoop bound_int n loop+    pat' <- renamePat pat+    renameForLoop loop $ \_loop' val_pats' _form' i' it' _bound' loop_vars' body' -> do+      remain_body <- insertStmsM $ do+        i_remain <-+          letExp "i_remain" $+            BasicOp $ BinOp (Add it OverflowUndef) total_iters (Var i')+        pure $ substituteNames (M.singleton i' i_remain) body'+      let loop_params_rem = map fst val_pats'+          loop_inits_rem = map (Var . patElemName) $ patElems pat'+          val_pats_rem = zip loop_params_rem loop_inits_rem+          remain_loop = DoLoop val_pats_rem (ForLoop i' it' remain_iters loop_vars') remain_body+      collectStms_ $ do+        letBind pat' mined_loop+        letBind pat remain_loop++-- | Stripmines a statement. Only has an effect when the statement's+-- expression is a for-loop with a @#[stripmine(n)]@ attribute, where+-- @n@ is the nesting depth.+stripmineStm :: Stm SOACS -> ADM (Stms SOACS)+stripmineStm stm@(Let pat aux loop@(DoLoop _ ForLoop {} _)) =+  case nums of+    (n : _) -> stripmine n pat loop+    _ -> pure $ oneStm stm+  where+    extractNum (AttrComp "stripmine" [AttrInt n]) = Just n+    extractNum _ = Nothing+    nums = catMaybes $ mapAttrs extractNum $ stmAuxAttrs aux+stripmineStm stm = pure $ oneStm stm++stripmineStms :: Stms SOACS -> ADM (Stms SOACS)+stripmineStms = traverseFold stripmineStm++-- | Forward pass transformation of a loop. This includes modifying the loop+-- to save the loop values at each iteration onto a tape as well as copying+-- any consumed arrays in the loop's body and consuming said copies in lieu of+-- the originals (which will be consumed later in the reverse pass).+fwdLoop :: Pat Type -> StmAux () -> Exp SOACS -> ADM ()+fwdLoop pat aux loop =+  bindForLoop loop $ \val_pats form i _it bound _loop_vars body -> do+    bound64 <- asIntS Int64 bound+    let loop_params = map fst val_pats+        is_true_dep = inAttrs (AttrName "true_dep") . paramAttrs+        dont_copy_params = filter is_true_dep loop_params+        dont_copy = map paramName dont_copy_params+        loop_params_to_copy = loop_params \\ dont_copy_params++    empty_saved_array <-+      forM loop_params_to_copy $ \p ->+        letSubExp (baseString (paramName p) <> "_empty_saved")+          =<< eBlank (arrayOf (paramDec p) (Shape [bound64]) NoUniqueness)++    (body', (saved_pats, saved_params)) <- buildBody $+      localScope (scopeOfFParams loop_params) $+        inScopeOf form $ do+          copy_substs <- copyConsumedArrsInBody dont_copy body+          addStms $ bodyStms body+          i_i64 <- asIntS Int64 $ Var i+          (saved_updates, saved_pats_params) <- fmap unzip $+            forM loop_params_to_copy $ \p -> do+              let v = paramName p+                  t = paramDec p+              saved_param_v <- newVName $ baseString v <> "_saved"+              saved_pat_v <- newVName $ baseString v <> "_saved"+              setLoopTape v saved_pat_v+              let saved_param = Param mempty saved_param_v $ arrayOf t (Shape [bound64]) Unique+                  saved_pat = PatElem saved_pat_v $ arrayOf t (Shape [bound64]) NoUniqueness+              saved_update <-+                localScope (scopeOfFParams [saved_param]) $+                  letInPlace+                    (baseString v <> "_saved_update")+                    saved_param_v+                    (fullSlice (fromDecl $ paramDec saved_param) [DimFix i_i64])+                    $ substituteNames copy_substs $ BasicOp $ SubExp $ Var v+              pure (saved_update, (saved_pat, saved_param))+          pure (bodyResult body <> varsRes saved_updates, unzip saved_pats_params)++    let pat' = pat <> Pat saved_pats+        val_pats' = val_pats <> zip saved_params empty_saved_array+    addStm $ Let pat' aux $ DoLoop val_pats' form body'++-- | Construct a loop value-pattern for the adjoint of the+-- given variable.+valPatAdj :: VName -> ADM (Param DeclType, SubExp)+valPatAdj v = do+  v_adj <- adjVName v+  init_adj <- lookupAdjVal v+  t <- lookupType init_adj+  pure (Param mempty v_adj (toDecl t Unique), Var init_adj)++valPatAdjs :: LoopInfo [VName] -> ADM (LoopInfo [(Param DeclType, SubExp)])+valPatAdjs = (mapM . mapM) valPatAdj++-- | Reverses a loop by substituting the loop index as well as reversing+-- the arrays that loop variables are bound to.+reverseIndices :: Exp SOACS -> ADM (Substitutions, Substitutions, Stms SOACS)+reverseIndices loop = do+  bindForLoop loop $ \_val_pats form i it bound loop_vars _body -> do+    bound_minus_one <-+      inScopeOf form $+        let one = Constant $ IntValue $ intValue it (1 :: Int)+         in letSubExp "bound-1" $ BasicOp $ BinOp (Sub it OverflowUndef) bound one++    var_arrays_substs <- fmap M.fromList $+      inScopeOf form $ do+        forM (map snd loop_vars) $ \xs -> do+          xs_t <- lookupType xs+          xs_rev <-+            letExp "reverse" $+              BasicOp $+                Index xs $+                  fullSlice+                    xs_t+                    [DimSlice bound_minus_one bound (Constant (IntValue (Int64Value (-1))))]+          pure (xs, xs_rev)++    (i_rev, i_stms) <- collectStms $+      inScopeOf form $ do+        letExp (baseString i <> "_rev") $+          BasicOp $ BinOp (Sub it OverflowWrap) bound_minus_one (Var i)++    pure (var_arrays_substs, M.singleton i i_rev, i_stms)++-- | Pures a substitution which substitutes values in the reverse+-- loop body with values from the tape.+restore :: Stms SOACS -> [Param DeclType] -> VName -> ADM Substitutions+restore stms_adj loop_params' i' =+  M.fromList . catMaybes <$> mapM f loop_params'+  where+    dont_copy =+      map paramName $ filter (inAttrs (AttrName "true_dep") . paramAttrs) loop_params'+    f p+      | v `notElem` dont_copy = do+          m_vs <- lookupLoopTape v+          case m_vs of+            Nothing -> pure Nothing+            Just vs -> do+              vs_t <- lookupType vs+              i_i64' <- asIntS Int64 $ Var i'+              v' <- letExp "restore" $ BasicOp $ Index vs $ fullSlice vs_t [DimFix i_i64']+              t <- lookupType v+              v'' <- case (t, v `elem` consumed) of+                (Array {}, True) -> letExp "restore_copy" $ BasicOp $ Copy v'+                _ -> pure v'+              pure $ Just (v, v'')+      | otherwise = pure Nothing+      where+        v = paramName p+        consumed = namesToList $ consumedInStms $ fst $ Alias.analyseStms mempty stms_adj++-- | A type to keep track of and seperate values corresponding to different+-- parts of the loop.+data LoopInfo a = LoopInfo+  { loopRes :: a,+    loopFree :: a,+    loopVars :: a,+    loopVals :: a+  }+  deriving (Functor, Foldable, Traversable, Show)++-- | Transforms a for-loop into its reverse-mode derivative.+revLoop :: (Stms SOACS -> ADM ()) -> Pat Type -> Exp SOACS -> ADM ()+revLoop diffStms pat loop =+  bindForLoop loop $ \val_pats _form _i _it _bound _loop_vars _body ->+    renameForLoop loop $+      \loop' val_pats' form' i' _it' _bound' loop_vars' body' -> do+        let loop_params = map fst val_pats+            (loop_params', loop_vals') = unzip val_pats'+            loop_var_arrays' = map snd loop_vars'+            getVName Constant {} = Nothing+            getVName (Var v) = Just v+            loop_vnames =+              LoopInfo+                { loopRes = mapMaybe subExpResVName $ bodyResult body',+                  loopFree =+                    (namesToList (freeIn loop') \\ loop_var_arrays') \\ mapMaybe getVName loop_vals',+                  loopVars = loop_var_arrays',+                  loopVals = nub $ mapMaybe getVName loop_vals'+                }++        renameLoopTape $ M.fromList $ zip (map paramName loop_params) (map paramName loop_params')++        forM_ (zip (bodyResult body') $ patElems pat) $ \(se_res, pe) ->+          case subExpResVName se_res of+            Just v -> setAdj v =<< lookupAdj (patElemName pe)+            Nothing -> pure ()++        (var_array_substs, i_subst, i_stms) <-+          reverseIndices loop'++        val_pat_adjs <- valPatAdjs loop_vnames+        let val_pat_adjs_list = concat $ toList val_pat_adjs++        (loop_adjs, stms_adj) <- collectStms $+          inScopeOf form' $+            localScope (scopeOfFParams (map fst val_pat_adjs_list <> loop_params')) $ do+              addStms i_stms+              (loop_adjs, stms_adj) <- collectStms $+                subAD $ do+                  zipWithM_+                    (\val_pat v -> insAdj v (paramName $ fst val_pat))+                    val_pat_adjs_list+                    (concat $ toList loop_vnames)+                  diffStms $ bodyStms body'++                  let update_var_arrays v vs = do+                        vs_t <- lookupType vs+                        v_adj <- lookupAdjVal v+                        updateAdjSlice (fullSlice vs_t [DimFix $ Var i']) vs v_adj+                  zipWithM_+                    update_var_arrays+                    (map (paramName . fst) loop_vars')+                    (loopVars loop_vnames)++                  loop_res_adjs <- mapM (lookupAdjVal . paramName) loop_params'+                  loop_free_adjs <- mapM lookupAdjVal $ loopFree loop_vnames+                  loop_vars_adjs <- mapM lookupAdjVal $ loopVars loop_vnames+                  loop_vals_adjs <- mapM lookupAdjVal $ loopVals loop_vnames++                  pure $+                    LoopInfo+                      { loopRes = loop_res_adjs,+                        loopFree = loop_free_adjs,+                        loopVars = loop_vars_adjs,+                        loopVals = loop_vals_adjs+                      }+              (substs, restore_stms) <-+                collectStms $ restore stms_adj loop_params' i'+              addStms $ substituteNames i_subst restore_stms+              addStms $ substituteNames i_subst $ substituteNames substs stms_adj+              pure loop_adjs++        inScopeOf stms_adj $+          localScope (scopeOfFParams $ map fst val_pat_adjs_list) $ do+            let body_adj = mkBody stms_adj $ varsRes $ concat $ toList loop_adjs+                restore_true_deps = M.fromList $+                  flip mapMaybe (zip loop_params' $ patElems pat) $ \(p, pe) ->+                    if p `elem` filter (inAttrs (AttrName "true_dep") . paramAttrs) loop_params'+                      then Just (paramName p, patElemName pe)+                      else Nothing+            adjs' <-+              letTupExp "loop_adj" $+                substituteNames (restore_true_deps <> var_array_substs) $+                  DoLoop val_pat_adjs_list form' body_adj+            let (loop_res_adjs, loop_free_var_val_adjs) =+                  splitAt (length $ loopRes loop_adjs) adjs'+                (loop_free_adjs, loop_var_val_adjs) =+                  splitAt (length $ loopFree loop_adjs) loop_free_var_val_adjs+                (loop_var_adjs, loop_val_adjs) =+                  splitAt (length $ loopVars loop_adjs) loop_var_val_adjs+            returnSweepCode $ do+              zipWithM_ updateSubExpAdj loop_vals' loop_res_adjs+              zipWithM_ insAdj (loopFree loop_vnames) loop_free_adjs+              zipWithM_ insAdj (loopVars loop_vnames) loop_var_adjs+              zipWithM_ updateAdj (loopVals loop_vnames) loop_val_adjs++-- | Transforms a loop into its reverse-mode derivative.+diffLoop :: (Stms SOACS -> ADM ()) -> Pat Type -> StmAux () -> Exp SOACS -> ADM () -> ADM ()+diffLoop diffStms pat aux loop m+  | isWhileLoop loop =+      let getBound (AttrComp "bound" [AttrInt b]) = Just b+          getBound _ = Nothing+          bounds = catMaybes $ mapAttrs getBound $ stmAuxAttrs aux+       in case bounds of+            (bound : _) -> do+              let bound_se = Constant $ IntValue $ intValue Int64 bound+              for_loop <- convertWhileLoop bound_se loop+              diffLoop diffStms pat aux for_loop m+            _ -> do+              bound <- computeWhileIters loop+              for_loop <- convertWhileLoop bound =<< renameExp loop+              diffLoop diffStms pat aux for_loop m+  | otherwise = do+      fwdLoop pat aux loop+      m+      revLoop diffStms pat loop

@@ -0,0 +1,233 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeFamilies #-}++module Futhark.AD.Rev.Map (vjpMap) where++import Control.Monad+import Data.Bifunctor (first)+import Futhark.AD.Rev.Monad+import Futhark.Analysis.PrimExp.Convert+import Futhark.Builder+import Futhark.IR.SOACS+import Futhark.Tools+import Futhark.Transform.Rename+import Futhark.Util (splitAt3)++-- | A classification of a free variable based on its adjoint.  The+-- 'VName' stored is *not* the adjoint, but the primal variable.+data AdjVar+  = -- | Adjoint is already an accumulator.+    FreeAcc VName+  | -- | Currently has no adjoint, but should be given one, and is an+    -- array with this shape and element type.+    FreeArr VName Shape PrimType+  | -- | Does not need an accumulator adjoint (might still be an array).+    FreeNonAcc VName++classifyAdjVars :: [VName] -> ADM [AdjVar]+classifyAdjVars = mapM f+  where+    f v = do+      v_adj <- lookupAdjVal v+      v_adj_t <- lookupType v_adj+      case v_adj_t of+        Array pt shape _ ->+          pure $ FreeArr v shape pt+        Acc {} ->+          pure $ FreeAcc v+        _ ->+          pure $ FreeNonAcc v++partitionAdjVars :: [AdjVar] -> ([(VName, (Shape, PrimType))], [VName], [VName])+partitionAdjVars [] = ([], [], [])+partitionAdjVars (fv : fvs) =+  case fv of+    FreeArr v shape t -> ((v, (shape, t)) : xs, ys, zs)+    FreeAcc v -> (xs, v : ys, zs)+    FreeNonAcc v -> (xs, ys, v : zs)+  where+    (xs, ys, zs) = partitionAdjVars fvs++buildRenamedBody ::+  MonadBuilder m =>+  m (Result, a) ->+  m (Body (Rep m), a)+buildRenamedBody m = do+  (body, x) <- buildBody m+  body' <- renameBody body+  pure (body', x)++withAcc ::+  [(Shape, [VName], Maybe (Lambda SOACS, [SubExp]))] ->+  ([VName] -> ADM Result) ->+  ADM [VName]+withAcc [] m =+  mapM (letExp "withacc_res" . BasicOp . SubExp . resSubExp) =<< m []+withAcc inputs m = do+  (cert_params, acc_params) <- fmap unzip $+    forM inputs $ \(shape, arrs, _) -> do+      cert_param <- newParam "acc_cert_p" $ Prim Unit+      ts <- mapM (fmap (stripArray (shapeRank shape)) . lookupType) arrs+      acc_param <- newParam "acc_p" $ Acc (paramName cert_param) shape ts NoUniqueness+      pure (cert_param, acc_param)+  acc_lam <-+    subAD $ mkLambda (cert_params ++ acc_params) $ m $ map paramName acc_params+  letTupExp "withhacc_res" $ WithAcc inputs acc_lam++vjpMap :: VjpOps -> [Adj] -> StmAux () -> SubExp -> Lambda SOACS -> [VName] -> ADM ()+vjpMap ops res_adjs _ w map_lam as+  | Just res_ivs <- mapM isSparse res_adjs = returnSweepCode $ do+      -- Since at most only a constant number of adjoint are nonzero+      -- (length res_ivs), there is no need for the return sweep code to+      -- contain a Map at all.++      free <- filterM isActive $ namesToList $ freeIn map_lam+      free_ts <- mapM lookupType free+      let adjs_for = map paramName (lambdaParams map_lam) ++ free+          adjs_ts = map paramType (lambdaParams map_lam) ++ free_ts++      let oneHot res_i adj_v = zipWith f [0 :: Int ..] $ lambdaReturnType map_lam+            where+              f j t+                | res_i == j = adj_v+                | otherwise = AdjZero (arrayShape t) (elemType t)+          -- Values for the out-of-bounds case does not matter, as we will+          -- be writing to an out-of-bounds index anyway, which is ignored.+          ooBounds adj_i = subAD . buildRenamedBody $ do+            forM_ (zip as adjs_ts) $ \(a, t) -> do+              scratch <- letSubExp "oo_scratch" =<< eBlank t+              updateAdjIndex a (OutOfBounds, adj_i) scratch+            first subExpsRes . adjsReps <$> mapM lookupAdj as+          inBounds res_i adj_i adj_v = subAD . buildRenamedBody $ do+            forM_ (zip (lambdaParams map_lam) as) $ \(p, a) -> do+              a_t <- lookupType a+              letBindNames [paramName p] $+                BasicOp $ Index a $ fullSlice a_t [DimFix adj_i]+            adj_elems <-+              fmap (map resSubExp) . bodyBind . lambdaBody+                =<< vjpLambda ops (oneHot res_i (AdjVal adj_v)) adjs_for map_lam+            forM_ (zip as adj_elems) $ \(a, a_adj_elem) -> do+              updateAdjIndex a (AssumeBounds, adj_i) a_adj_elem+            first subExpsRes . adjsReps <$> mapM lookupAdj as++          -- Generate an iteration of the map function for every+          -- position.  This is a bit inefficient - probably we could do+          -- some deduplication.+          forPos res_i (check, adj_i, adj_v) = do+            as_adj <-+              case check of+                CheckBounds b -> do+                  (obbranch, mkadjs) <- ooBounds adj_i+                  (ibbranch, _) <- inBounds res_i adj_i adj_v+                  fmap mkadjs . letTupExp' "map_adj_elem"+                    =<< eIf+                      (maybe (eDimInBounds (eSubExp w) (eSubExp adj_i)) eSubExp b)+                      (pure ibbranch)+                      (pure obbranch)+                AssumeBounds -> do+                  (body, mkadjs) <- inBounds res_i adj_i adj_v+                  mkadjs . map resSubExp <$> bodyBind body+                OutOfBounds ->+                  mapM lookupAdj as++            zipWithM setAdj as as_adj++          -- Generate an iteration of the map function for every result.+          forRes res_i = mapM_ (forPos res_i)++      zipWithM_ forRes [0 ..] res_ivs+  where+    isSparse (AdjSparse (Sparse shape _ ivs)) = do+      guard $ shapeDims shape == [w]+      Just ivs+    isSparse _ =+      Nothing+-- See Note [Adjoints of accumulators] for how we deal with+-- accumulators - it's a bit tricky here.+vjpMap ops pat_adj aux w map_lam as = returnSweepCode $ do+  pat_adj_vals <- forM (zip pat_adj (lambdaReturnType map_lam)) $ \(adj, t) ->+    case t of+      Acc {} -> letExp "acc_adj_rep" . BasicOp . Replicate (Shape [w]) . Var =<< adjVal adj+      _ -> adjVal adj+  pat_adj_params <-+    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)++    -- 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')+    forM_ (zip3 param_ts as param_contribs) $ \(param_t, a, param_contrib) ->+      case param_t of+        Acc {} -> freeContrib a param_contrib+        _ -> updateAdj a 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++    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]))++    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++    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 (baseString v <> "_contrib_sum") $+              Op $ Screma w [contribs] reduce+          void $ updateAdj v contrib_sum

@@ -0,0 +1,551 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GeneralizedNewtypeDeriving #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeFamilies #-}++-- Naming scheme:+--+-- An adjoint-related object for "x" is named "x_adj".  This means+-- both actual adjoints and statements.+--+-- Do not assume "x'" means anything related to derivatives.+module Futhark.AD.Rev.Monad+  ( ADM,+    RState (..),+    runADM,+    Adj (..),+    InBounds (..),+    Sparse (..),+    adjFromParam,+    adjFromVar,+    lookupAdj,+    lookupAdjVal,+    adjVal,+    updateAdj,+    updateSubExpAdj,+    updateAdjSlice,+    updateAdjIndex,+    setAdj,+    insAdj,+    insSubExpAdj,+    adjsReps,+    --+    copyConsumedArrsInStm,+    copyConsumedArrsInBody,+    addSubstitution,+    returnSweepCode,+    --+    adjVName,+    subAD,+    noAdjsFor,+    subSubsts,+    isActive,+    --+    tabNest,+    oneExp,+    zeroExp,+    unitAdjOfType,+    addLambda,+    --+    VjpOps (..),+    --+    setLoopTape,+    lookupLoopTape,+    substLoopTape,+    renameLoopTape,+  )+where++import Control.Monad+import Control.Monad.State.Strict+import Data.Bifunctor (second)+import Data.List (foldl')+import qualified Data.Map as M+import Data.Maybe+import qualified Futhark.Analysis.Alias as Alias+import Futhark.Analysis.PrimExp.Convert+import Futhark.Builder+import Futhark.IR.Aliases (consumedInStms)+import Futhark.IR.Prop.Aliases+import Futhark.IR.SOACS+import Futhark.Tools+import Futhark.Transform.Substitute+import Futhark.Util (chunks)++zeroExp :: Type -> Exp rep+zeroExp (Prim pt) =+  BasicOp $ SubExp $ Constant $ blankPrimValue pt+zeroExp (Array pt shape _) =+  BasicOp $ Replicate shape $ Constant $ blankPrimValue pt+zeroExp t = error $ "zeroExp: " ++ pretty t++onePrim :: PrimType -> PrimValue+onePrim (IntType it) = IntValue $ intValue it (1 :: Int)+onePrim (FloatType ft) = FloatValue $ floatValue ft (1 :: Double)+onePrim Bool = BoolValue True+onePrim Unit = UnitValue++oneExp :: Type -> Exp rep+oneExp (Prim t) = BasicOp $ SubExp $ constant $ onePrim t+oneExp (Array pt shape _) =+  BasicOp $ Replicate shape $ Constant $ onePrim pt+oneExp t = error $ "oneExp: " ++ pretty t++-- | Whether 'Sparse' should check bounds or assume they are correct.+-- The latter results in simpler code.+data InBounds+  = -- | If a SubExp is provided, it references a boolean that is true+    -- when in-bounds.+    CheckBounds (Maybe SubExp)+  | AssumeBounds+  | -- | Dynamically these will always fail, so don't bother+    -- generating code for the update.  This is only needed to ensure+    -- a consistent representation of sparse Jacobians.+    OutOfBounds+  deriving (Eq, Ord, Show)++-- | A symbolic representation of an array that is all zeroes, except at one+-- index.+data Sparse = Sparse+  { -- | The shape of the array.+    sparseShape :: Shape,+    -- | Element type of the array.+    sparseType :: PrimType,+    -- | Locations and values of nonzero values.  Indexes may be+    -- negative, in which case the value is ignored (unless+    -- 'AssumeBounds' is used).+    sparseIdxVals :: [(InBounds, SubExp, SubExp)]+  }+  deriving (Eq, Ord, Show)++-- | The adjoint of a variable.+data Adj+  = AdjSparse Sparse+  | AdjVal SubExp+  | AdjZero Shape PrimType+  deriving (Eq, Ord, Show)++instance Substitute Adj where+  substituteNames m (AdjVal (Var v)) = AdjVal $ Var $ substituteNames m v+  substituteNames _ adj = adj++zeroArray :: MonadBuilder m => Shape -> Type -> m VName+zeroArray shape t+  | shapeRank shape == 0 =+      letExp "zero" $ zeroExp t+  | otherwise = do+      zero <- letSubExp "zero" $ zeroExp t+      attributing (oneAttr "sequential") $+        letExp "zeroes_" $ BasicOp $ Replicate shape zero++sparseArray :: (MonadBuilder m, Rep m ~ SOACS) => Sparse -> m VName+sparseArray (Sparse shape t ivs) = do+  flip (foldM f) ivs =<< zeroArray shape (Prim t)+  where+    arr_t = Prim t `arrayOfShape` shape+    f arr (check, i, se) = do+      let stm s =+            letExp "sparse" . BasicOp $+              Update s arr (fullSlice arr_t [DimFix i]) se+      case check of+        AssumeBounds -> stm Unsafe+        CheckBounds _ -> stm Safe+        OutOfBounds -> pure arr++adjFromVar :: VName -> Adj+adjFromVar = AdjVal . Var++adjFromParam :: Param t -> Adj+adjFromParam = adjFromVar . paramName++unitAdjOfType :: Type -> ADM Adj+unitAdjOfType t = AdjVal <$> letSubExp "adj_unit" (oneExp t)++-- | The values representing an adjoint in symbolic form.  This is+-- used for when we wish to return an Adj from a Body or similar+-- without forcing manifestation.  Also returns a function for+-- reassembling the Adj from a new representation (the list must have+-- the same length).+adjRep :: Adj -> ([SubExp], [SubExp] -> Adj)+adjRep (AdjVal se) = ([se], \[se'] -> AdjVal se')+adjRep (AdjZero shape pt) = ([], \[] -> AdjZero shape pt)+adjRep (AdjSparse (Sparse shape pt ivs)) =+  (concatMap ivRep ivs, AdjSparse . Sparse shape pt . repIvs ivs)+  where+    ivRep (_, i, v) = [i, v]+    repIvs ((check, _, _) : ivs') (i : v : ses) =+      (check', i, v) : repIvs ivs' ses+      where+        check' = case check of+          AssumeBounds -> AssumeBounds+          CheckBounds b -> CheckBounds b+          OutOfBounds -> CheckBounds (Just (constant False)) -- sic!+    repIvs _ _ = []++-- | Conveniently convert a list of Adjs to their representation, as+-- well as produce a function for converting back.+adjsReps :: [Adj] -> ([SubExp], [SubExp] -> [Adj])+adjsReps adjs =+  let (reps, fs) = unzip $ map adjRep adjs+   in (concat reps, zipWith ($) fs . chunks (map length reps))++data RState = RState+  { stateAdjs :: M.Map VName Adj,+    stateLoopTape :: Substitutions,+    stateSubsts :: Substitutions,+    stateNameSource :: VNameSource+  }++newtype ADM a = ADM (BuilderT SOACS (State RState) a)+  deriving+    ( Functor,+      Applicative,+      Monad,+      MonadState RState,+      MonadFreshNames,+      HasScope SOACS,+      LocalScope SOACS+    )++instance MonadBuilder ADM where+  type Rep ADM = SOACS+  mkExpDecM pat e = ADM $ mkExpDecM pat e+  mkBodyM bnds res = ADM $ mkBodyM bnds res+  mkLetNamesM pat e = ADM $ mkLetNamesM pat e++  addStms = ADM . addStms+  collectStms (ADM m) = ADM $ collectStms m++instance MonadFreshNames (State RState) where+  getNameSource = gets stateNameSource+  putNameSource src = modify (\env -> env {stateNameSource = src})++runADM :: MonadFreshNames m => ADM a -> m a+runADM (ADM m) =+  modifyNameSource $ \vn ->+    second stateNameSource $+      runState+        (fst <$> runBuilderT m mempty)+        (RState mempty mempty mempty vn)++adjVal :: Adj -> ADM VName+adjVal (AdjVal se) = letExp "const_adj" $ BasicOp $ SubExp se+adjVal (AdjSparse sparse) = sparseArray sparse+adjVal (AdjZero shape t) = zeroArray shape $ Prim t++setAdj :: VName -> Adj -> ADM ()+setAdj v v_adj = modify $ \env ->+  env {stateAdjs = M.insert v v_adj $ stateAdjs env}++insAdj :: VName -> VName -> ADM ()+insAdj v = setAdj v . AdjVal . Var++adjVName :: VName -> ADM VName+adjVName v = newVName (baseString v <> "_adj")++-- | Create copies of all arrays consumed in the given statement, and+-- return statements which include copies of the consumed arrays.+--+-- See Note [Consumption].+copyConsumedArrsInStm :: Stm SOACS -> ADM (Substitutions, Stms SOACS)+copyConsumedArrsInStm s = inScopeOf s $ collectStms $ copyConsumedArrsInStm' s+  where+    copyConsumedArrsInStm' stm =+      let onConsumed v = inScopeOf s $ do+            v_t <- lookupType v+            case v_t of+              Array {} -> do+                v' <- letExp (baseString v <> "_ad_copy") (BasicOp $ Copy v)+                addSubstitution v' v+                pure [(v, v')]+              _ -> pure mempty+       in M.fromList . mconcat+            <$> mapM onConsumed (namesToList $ consumedInStms $ fst (Alias.analyseStms mempty (oneStm stm)))++copyConsumedArrsInBody :: [VName] -> Body SOACS -> ADM Substitutions+copyConsumedArrsInBody dontCopy b =+  mconcat <$> mapM onConsumed (filter (`notElem` dontCopy) $ namesToList $ consumedInBody (Alias.analyseBody mempty b))+  where+    onConsumed v = do+      v_t <- lookupType v+      case v_t of+        Acc {} -> error $ "copyConsumedArrsInBody: Acc " <> pretty v+        Array {} -> M.singleton v <$> letExp (baseString v <> "_ad_copy") (BasicOp $ Copy v)+        _ -> pure mempty++returnSweepCode :: ADM a -> ADM a+returnSweepCode m = do+  (a, stms) <- collectStms m+  substs <- gets stateSubsts+  addStms $ substituteNames substs stms+  pure a++addSubstitution :: VName -> VName -> ADM ()+addSubstitution v v' = modify $ \env ->+  env {stateSubsts = M.insert v v' $ stateSubsts env}++-- While evaluating this action, pretend these variables have no+-- adjoints.  Restore current adjoints afterwards.  This is used for+-- handling certain nested operations. XXX: feels like this should+-- really be part of subAD, somehow.  Main challenge is that we don't+-- want to blank out Accumulator adjoints.  Also, might be inefficient+-- to blank out array adjoints.+noAdjsFor :: Names -> ADM a -> ADM a+noAdjsFor names m = do+  old <- gets $ \env -> mapMaybe (`M.lookup` stateAdjs env) names'+  modify $ \env -> env {stateAdjs = foldl' (flip M.delete) (stateAdjs env) names'}+  x <- m+  modify $ \env -> env {stateAdjs = M.fromList (zip names' old) <> stateAdjs env}+  pure x+  where+    names' = namesToList names++addBinOp :: PrimType -> BinOp+addBinOp (IntType it) = Add it OverflowWrap+addBinOp (FloatType ft) = FAdd ft+addBinOp Bool = LogAnd+addBinOp Unit = LogAnd++tabNest :: Int -> [VName] -> ([VName] -> [VName] -> ADM [VName]) -> ADM [VName]+tabNest = tabNest' []+  where+    tabNest' is 0 vs f = f (reverse is) vs+    tabNest' is n vs f = do+      vs_ts <- mapM lookupType vs+      let w = arraysSize 0 vs_ts+      iota <-+        letExp "tab_iota" . BasicOp $+          Iota w (intConst Int64 0) (intConst Int64 1) Int64+      iparam <- newParam "i" $ Prim int64+      params <- forM vs $ \v ->+        newParam (baseString v <> "_p") . rowType =<< lookupType v+      ((ret, res), stms) <- collectStms . localScope (scopeOfLParams (iparam : params)) $ do+        res <- tabNest' (paramName iparam : is) (n - 1) (map paramName params) f+        ret <- mapM lookupType res+        pure (ret, varsRes res)+      let lam = Lambda (iparam : params) (Body () stms res) ret+      letTupExp "tab" $ Op $ Screma w (iota : vs) (mapSOAC lam)++-- | Construct a lambda for adding two values of the given type.+addLambda :: Type -> ADM (Lambda SOACS)+addLambda (Prim pt) = binOpLambda (addBinOp pt) pt+addLambda t@Array {} = do+  xs_p <- newParam "xs" t+  ys_p <- newParam "ys" t+  lam <- addLambda $ rowType t+  body <- insertStmsM $ do+    res <-+      letSubExp "lam_map" . Op $+        Screma (arraySize 0 t) [paramName xs_p, paramName ys_p] (mapSOAC lam)+    pure $ resultBody [res]+  pure+    Lambda+      { lambdaParams = [xs_p, ys_p],+        lambdaReturnType = [t],+        lambdaBody = body+      }+addLambda t =+  error $ "addLambda: " ++ show t++-- Construct an expression for adding the two variables.+addExp :: VName -> VName -> ADM (Exp SOACS)+addExp x y = do+  x_t <- lookupType x+  case x_t of+    Prim pt ->+      pure $ BasicOp $ BinOp (addBinOp pt) (Var x) (Var y)+    Array {} -> do+      lam <- addLambda $ rowType x_t+      pure $ Op $ Screma (arraySize 0 x_t) [x, y] (mapSOAC lam)+    _ ->+      error $ "addExp: unexpected type: " ++ pretty x_t++lookupAdj :: VName -> ADM Adj+lookupAdj v = do+  maybeAdj <- gets $ M.lookup v . stateAdjs+  case maybeAdj of+    Nothing -> do+      v_t <- lookupType v+      case v_t of+        Acc _ shape [Prim t] _ -> pure $ AdjZero shape t+        _ -> pure $ AdjZero (arrayShape v_t) (elemType v_t)+    Just v_adj -> pure v_adj++lookupAdjVal :: VName -> ADM VName+lookupAdjVal v = adjVal =<< lookupAdj v++updateAdj :: VName -> VName -> ADM ()+updateAdj v d = do+  maybeAdj <- gets $ M.lookup v . stateAdjs+  case maybeAdj of+    Nothing ->+      insAdj v d+    Just adj -> do+      v_adj <- adjVal adj+      v_adj_t <- lookupType v_adj+      case v_adj_t of+        Acc {} -> do+          dims <- arrayDims <$> lookupType d+          ~[v_adj'] <-+            tabNest (length dims) [d, v_adj] $ \is [d', v_adj'] ->+              letTupExp "acc" $+                BasicOp $ UpdateAcc v_adj' (map Var is) [Var d']+          insAdj v v_adj'+        _ -> do+          v_adj' <- letExp (baseString v <> "_adj") =<< addExp v_adj d+          insAdj v v_adj'++updateAdjSlice :: Slice SubExp -> VName -> VName -> ADM ()+updateAdjSlice (Slice [DimFix i]) v d =+  updateAdjIndex v (AssumeBounds, i) (Var d)+updateAdjSlice slice v d = do+  t <- lookupType v+  v_adj <- lookupAdjVal v+  v_adj_t <- lookupType v_adj+  v_adj' <- case v_adj_t of+    Acc {} -> do+      let dims = sliceDims slice+      ~[v_adj'] <-+        tabNest (length dims) [d, v_adj] $ \is [d', v_adj'] -> do+          slice' <-+            traverse (toSubExp "index") $+              fixSlice (fmap pe64 slice) $ map le64 is+          letTupExp (baseString v_adj') . BasicOp $+            UpdateAcc v_adj' slice' [Var d']+      pure v_adj'+    _ -> do+      v_adjslice <-+        if primType t+          then pure v_adj+          else letExp (baseString v ++ "_slice") $ BasicOp $ Index v_adj slice+      letInPlace "updated_adj" v_adj slice =<< addExp v_adjslice d+  insAdj v v_adj'++updateSubExpAdj :: SubExp -> VName -> ADM ()+updateSubExpAdj Constant {} _ = pure ()+updateSubExpAdj (Var v) d = void $ updateAdj v d++insSubExpAdj :: SubExp -> VName -> ADM ()+insSubExpAdj Constant {} _ = pure ()+insSubExpAdj (Var v) d = void $ insAdj v d++-- The index may be negative, in which case the update has no effect.+updateAdjIndex :: VName -> (InBounds, SubExp) -> SubExp -> ADM ()+updateAdjIndex v (check, i) se = do+  maybeAdj <- gets $ M.lookup v . stateAdjs+  t <- lookupType v+  let iv = (check, i, se)+  case maybeAdj of+    Nothing -> do+      setAdj v $ AdjSparse $ Sparse (arrayShape t) (elemType t) [iv]+    Just AdjZero {} ->+      setAdj v $ AdjSparse $ Sparse (arrayShape t) (elemType t) [iv]+    Just (AdjSparse (Sparse shape pt ivs)) ->+      setAdj v $ AdjSparse $ Sparse shape pt $ iv : ivs+    Just adj@AdjVal {} -> do+      v_adj <- adjVal adj+      v_adj_t <- lookupType v_adj+      se_v <- letExp "se_v" $ BasicOp $ SubExp se+      insAdj v+        =<< case v_adj_t of+          Acc {}+            | check == OutOfBounds ->+                pure v_adj+            | otherwise -> do+                dims <- arrayDims <$> lookupType se_v+                ~[v_adj'] <-+                  tabNest (length dims) [se_v, v_adj] $ \is [se_v', v_adj'] ->+                    letTupExp "acc" $+                      BasicOp $ UpdateAcc v_adj' (i : map Var is) [Var se_v']+                pure v_adj'+          _ -> do+            let stms s = do+                  v_adj_i <- letExp (baseString v_adj <> "_i") $ BasicOp $ Index v_adj $ fullSlice v_adj_t [DimFix i]+                  se_update <- letSubExp "updated_adj_i" =<< addExp se_v v_adj_i+                  letExp (baseString v_adj) $+                    BasicOp $ Update s v_adj (fullSlice v_adj_t [DimFix i]) se_update+            case check of+              CheckBounds _ -> stms Safe+              AssumeBounds -> stms Unsafe+              OutOfBounds -> pure v_adj++-- | Is this primal variable active in the AD sense?  FIXME: this is+-- (obviously) much too conservative.+isActive :: VName -> ADM Bool+isActive = fmap (/= Prim Unit) . lookupType++subAD :: ADM a -> ADM a+subAD m = do+  old_state_adjs <- gets stateAdjs+  x <- m+  modify $ \s -> s {stateAdjs = old_state_adjs}+  pure x++subSubsts :: ADM a -> ADM a+subSubsts m = do+  old_state_substs <- gets stateSubsts+  x <- m+  modify $ \s -> s {stateSubsts = old_state_substs}+  pure x++data VjpOps = VjpOps+  { vjpLambda :: [Adj] -> [VName] -> Lambda SOACS -> ADM (Lambda SOACS),+    vjpStm :: Stm SOACS -> ADM () -> ADM ()+  }++-- | @setLoopTape v vs@ establishes @vs@ as the name of the array+-- where values of loop parameter @v@ from the forward pass are+-- stored.+setLoopTape :: VName -> VName -> ADM ()+setLoopTape v vs = modify $ \env ->+  env {stateLoopTape = M.insert v vs $ stateLoopTape env}++-- | Look-up the name of the array where @v@ is stored.+lookupLoopTape :: VName -> ADM (Maybe VName)+lookupLoopTape v = gets $ M.lookup v . stateLoopTape++-- | @substLoopTape v v'@ substitutes the key @v@ for @v'@. That is,+-- if @v |-> vs@ then after the substitution @v' |-> vs@ (and @v@+-- points to nothing).+substLoopTape :: VName -> VName -> ADM ()+substLoopTape v v' = mapM_ (setLoopTape v') =<< lookupLoopTape v++-- | Renames the keys of the loop tape. Useful for fixing the+-- the names in the loop tape after a loop rename.+renameLoopTape :: Substitutions -> ADM ()+renameLoopTape = mapM_ (uncurry substLoopTape) . M.toList++-- Note [Consumption]+--+-- Parts of this transformation depends on duplicating computation.+-- This is a problem when a primal expression consumes arrays (via+-- e.g. Update).  For example, consider how we handle this conditional:+--+--   if b then ys with [0] = 0 else ys+--+-- This consumes the array 'ys', which means that when we later+-- generate code for the return sweep, we can no longer use 'ys'.+-- This is a problem, because when we call 'diffBody' on the branch+-- bodies, we'll keep the primal code (maybe it'll be removed by+-- simplification later - we cannot know).  A similar issue occurs for+-- SOACs.  Our solution is to make copies of all consumes arrays:+--+--  let ys_copy = copy ys+--+-- Then we generate code for the return sweep as normal, but replace+-- _every instance_ of 'ys' in the generated code with 'ys_copy'.+-- This works because Futhark does not have *semantic* in-place+-- updates - any uniqueness violation can be replaced with copies (on+-- arrays, anyway).+--+-- If we are lucky, the uses of 'ys_copy' will be removed by+-- simplification, and there will be no overhead.  But even if not,+-- this is still (asymptotically) efficient because the array that is+-- being consumed must in any case have been produced within the code+-- that we are differentiating, so a copy is at most a scalar+-- overhead.  This is _not_ the case when loops are involved.+--+-- Also, the above only works for arrays, not accumulator variables.+-- Those will need some other mechanism.

@@ -0,0 +1,204 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeFamilies #-}++module Futhark.AD.Rev.Reduce+  ( diffReduce,+    diffMinMaxReduce,+  )+where++import Control.Monad+import Futhark.AD.Rev.Monad+import Futhark.Analysis.PrimExp.Convert+import Futhark.Builder+import Futhark.IR.SOACS+import Futhark.Tools+import Futhark.Transform.Rename++eReverse :: MonadBuilder m => VName -> m VName+eReverse arr = do+  arr_t <- lookupType arr+  let w = arraySize 0 arr_t+  start <-+    letSubExp "rev_start" $+      BasicOp $ BinOp (Sub Int64 OverflowUndef) w (intConst Int64 1)+  let stride = intConst Int64 (-1)+      slice = fullSlice arr_t [DimSlice start w stride]+  letExp (baseString arr <> "_rev") $ BasicOp $ Index arr slice++eRotate :: MonadBuilder m => [SubExp] -> VName -> m VName+eRotate rots arr = letExp (baseString arr <> "_rot") $ BasicOp $ Rotate rots arr++scanExc ::+  (MonadBuilder m, Rep m ~ SOACS) =>+  String ->+  Scan SOACS ->+  [VName] ->+  m [VName]+scanExc desc scan arrs = do+  w <- arraysSize 0 <$> mapM lookupType arrs+  form <- scanSOAC [scan]+  res_incl <- letTupExp (desc <> "_incl") $ Op $ Screma w arrs form+  res_incl_rot <- mapM (eRotate [intConst Int64 (-1)]) res_incl++  iota <-+    letExp "iota" . BasicOp $+      Iota w (intConst Int64 0) (intConst Int64 1) Int64++  iparam <- newParam "iota_param" $ Prim int64+  vparams <- mapM (newParam "vp") ts+  let params = iparam : vparams++  body <- runBodyBuilder . localScope (scopeOfLParams params) $ do+    let first_elem =+          eCmpOp+            (CmpEq int64)+            (eSubExp (Var (paramName iparam)))+            (eSubExp (intConst Int64 0))+    eBody+      [ eIf+          first_elem+          (resultBodyM nes)+          (resultBodyM $ map (Var . paramName) vparams)+      ]++  let lam = Lambda params body ts+  letTupExp desc $ Op $ Screma w (iota : res_incl_rot) (mapSOAC lam)+  where+    nes = scanNeutral scan+    ts = lambdaReturnType $ scanLambda scan++mkF :: Lambda SOACS -> ADM ([VName], Lambda SOACS)+mkF lam = do+  lam_l <- renameLambda lam+  lam_r <- renameLambda lam+  let q = length $ lambdaReturnType lam+      (lps, aps) = splitAt q $ lambdaParams lam_l+      (ips, rps) = splitAt q $ lambdaParams lam_r+  lam' <- mkLambda (lps <> aps <> rps) $ do+    lam_l_res <- bodyBind $ lambdaBody lam_l+    forM_ (zip ips lam_l_res) $ \(ip, SubExpRes cs se) ->+      certifying cs $ letBindNames [paramName ip] $ BasicOp $ SubExp se+    bodyBind $ lambdaBody lam_r+  pure (map paramName aps, lam')++diffReduce :: VjpOps -> [VName] -> SubExp -> [VName] -> Reduce SOACS -> ADM ()+diffReduce _ops [adj] w [a] red+  | Just [(op, _, _, _)] <- lamIsBinOp $ redLambda red,+    isAdd op = do+      adj_rep <-+        letExp (baseString adj <> "_rep") $+          BasicOp $ Replicate (Shape [w]) $ Var adj+      void $ updateAdj a adj_rep+  where+    isAdd FAdd {} = True+    isAdd Add {} = True+    isAdd _ = False+--+-- Differentiating a general single reduce:+--    let y = reduce \odot ne as+-- Forward sweep:+--    let ls = scan_exc \odot  ne as+--    let rs = scan_exc \odot' ne (reverse as)+-- Reverse sweep:+--    let as_c = map3 (f_bar y_bar) ls as (reverse rs)+-- where+--   x \odot' y = y \odot x+--   y_bar is the adjoint of the result y+--   f l_i a_i r_i = l_i \odot a_i \odot r_i+--   f_bar = the reverse diff of f with respect to a_i under the adjoint y_bar+-- The plan is to create+--   one scanomap SOAC which computes ls and rs+--   another map which computes as_c+--+diffReduce ops pat_adj w as red = do+  red' <- renameRed red+  flip_red <- renameRed =<< flipReduce red+  ls <- scanExc "ls" (redToScan red') as+  rs <-+    mapM eReverse+      =<< scanExc "ls" (redToScan flip_red)+      =<< mapM eReverse as++  (as_params, f) <- mkF $ redLambda red++  f_adj <- vjpLambda ops (map adjFromVar pat_adj) as_params f++  as_adj <- letTupExp "adjs" $ Op $ Screma w (ls ++ as ++ rs) (mapSOAC f_adj)++  zipWithM_ updateAdj as as_adj+  where+    renameRed (Reduce comm lam nes) =+      Reduce comm <$> renameLambda lam <*> pure nes++    redToScan :: Reduce SOACS -> Scan SOACS+    redToScan (Reduce _ lam nes) = Scan lam nes+    flipReduce (Reduce comm lam nes) = do+      lam' <- renameLambda lam {lambdaParams = flipParams $ lambdaParams lam}+      pure $ Reduce comm lam' nes+    flipParams ps = uncurry (flip (++)) $ splitAt (length ps `div` 2) ps++--+-- Special case of reduce with min/max:+--    let x = reduce minmax ne as+-- Forward trace (assuming w = length as):+--    let (x, x_ind) =+--      reduce (\ acc_v acc_i v i ->+--                 if (acc_v == v) then (acc_v, min acc_i i)+--                 else if (acc_v == minmax acc_v v)+--                      then (acc_v, acc_i)+--                      else (v, i))+--             (ne_min, -1)+--             (zip as (iota w))+-- Reverse trace:+--    num_elems = i64.bool (0 <= x_ind)+--    m_bar_repl = replicate num_elems m_bar+--    as_bar[x_ind:num_elems:1] += m_bar_repl+diffMinMaxReduce ::+  VjpOps -> VName -> StmAux () -> SubExp -> BinOp -> SubExp -> VName -> ADM () -> ADM ()+diffMinMaxReduce _ops x aux w minmax ne as m = do+  let t = binOpType minmax++  acc_v_p <- newParam "acc_v" $ Prim t+  acc_i_p <- newParam "acc_i" $ Prim int64+  v_p <- newParam "v" $ Prim t+  i_p <- newParam "i" $ Prim int64+  red_lam <-+    mkLambda [acc_v_p, acc_i_p, v_p, i_p] $+      fmap varsRes . letTupExp "idx_res"+        =<< eIf+          (eCmpOp (CmpEq t) (eParam acc_v_p) (eParam v_p))+          ( eBody+              [ eParam acc_v_p,+                eBinOp (SMin Int64) (eParam acc_i_p) (eParam i_p)+              ]+          )+          ( eBody+              [ eIf+                  ( eCmpOp+                      (CmpEq t)+                      (eParam acc_v_p)+                      (eBinOp minmax (eParam acc_v_p) (eParam v_p))+                  )+                  (eBody [eParam acc_v_p, eParam acc_i_p])+                  (eBody [eParam v_p, eParam i_p])+              ]+          )++  red_iota <-+    letExp "red_iota" $+      BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64+  form <- reduceSOAC [Reduce Commutative red_lam [ne, intConst Int64 (-1)]]+  x_ind <- newVName (baseString x <> "_ind")+  auxing aux $ letBindNames [x, x_ind] $ Op $ Screma w [as, red_iota] form++  m++  x_adj <- lookupAdjVal x+  in_bounds <-+    letSubExp "minmax_in_bounds" . BasicOp $+      CmpOp (CmpSlt Int64) (intConst Int64 0) w+  updateAdjIndex as (CheckBounds (Just in_bounds), Var x_ind) (Var x_adj)

@@ -0,0 +1,90 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeFamilies #-}++module Futhark.AD.Rev.SOAC (vjpSOAC) where++import Control.Monad+import Futhark.AD.Rev.Map+import Futhark.AD.Rev.Monad+import Futhark.AD.Rev.Reduce+import Futhark.AD.Rev.Scan+import Futhark.AD.Rev.Scatter+import Futhark.Analysis.PrimExp.Convert+import Futhark.Builder+import Futhark.IR.SOACS+import Futhark.Tools+import Futhark.Util (chunks)++-- We split any multi-op scan or reduction into multiple operations so+-- we can detect special cases.  Post-AD, the result may be fused+-- again.+splitScanRed ::+  VjpOps ->+  ([a] -> ADM (ScremaForm SOACS), a -> [SubExp]) ->+  (Pat Type, StmAux (), [a], SubExp, [VName]) ->+  ADM () ->+  ADM ()+splitScanRed vjpops (opSOAC, opNeutral) (pat, aux, ops, w, as) m = do+  let ks = map (length . opNeutral) ops+      pat_per_op = map Pat $ chunks ks $ patElems pat+      as_per_op = chunks ks as+      onOps (op : ops') (op_pat : op_pats') (op_as : op_as') = do+        op_form <- opSOAC [op]+        vjpSOAC vjpops op_pat aux (Screma w op_as op_form) $+          onOps ops' op_pats' op_as'+      onOps _ _ _ = m+  onOps ops pat_per_op as_per_op++commonSOAC :: Pat Type -> StmAux () -> SOAC SOACS -> ADM () -> ADM [Adj]+commonSOAC pat aux soac m = do+  addStm $ Let pat aux $ Op soac+  m+  returnSweepCode $ mapM lookupAdj $ patNames pat++vjpSOAC :: VjpOps -> Pat Type -> StmAux () -> SOAC SOACS -> ADM () -> ADM ()+vjpSOAC ops pat aux soac@(Screma w as form) m+  | Just reds <- isReduceSOAC form,+    length reds > 1 =+      splitScanRed ops (reduceSOAC, redNeutral) (pat, aux, reds, w, as) m+  | Just [red] <- isReduceSOAC form,+    [x] <- patNames pat,+    [ne] <- redNeutral red,+    [a] <- as,+    Just [(op, _, _, _)] <- lamIsBinOp $ redLambda red,+    isMinMaxOp op =+      diffMinMaxReduce ops x aux w op ne a m+  | Just red <- singleReduce <$> isReduceSOAC form = do+      pat_adj <- mapM adjVal =<< commonSOAC pat aux soac m+      diffReduce ops pat_adj w as red+  where+    isMinMaxOp (SMin _) = True+    isMinMaxOp (UMin _) = True+    isMinMaxOp (FMin _) = True+    isMinMaxOp (SMax _) = True+    isMinMaxOp (UMax _) = True+    isMinMaxOp (FMax _) = True+    isMinMaxOp _ = False+vjpSOAC ops pat aux soac@(Screma w as form) m+  | Just scans <- isScanSOAC form,+    length scans > 1 =+      splitScanRed ops (scanSOAC, scanNeutral) (pat, aux, scans, w, as) m+  | Just red <- singleScan <$> isScanSOAC form = do+      void $ commonSOAC pat aux soac m+      diffScan ops (patNames pat) w as red+vjpSOAC ops pat aux soac@(Screma w as form) m+  | Just lam <- isMapSOAC form = do+      pat_adj <- commonSOAC pat aux soac m+      vjpMap ops pat_adj aux w lam as+vjpSOAC ops pat _aux (Screma w as form) m+  | Just (reds, map_lam) <-+      isRedomapSOAC form = do+      (mapstm, redstm) <-+        redomapToMapAndReduce pat (w, reds, map_lam, as)+      vjpStm ops mapstm $ vjpStm ops redstm m+vjpSOAC ops pat aux (Scatter w lam ass written_info) m =+  vjpScatter ops pat aux (w, lam, ass, written_info) m+vjpSOAC _ _ _ soac _ =+  error $ "vjpSOAC unhandled:\n" ++ pretty soac

@@ -0,0 +1,154 @@+module Futhark.AD.Rev.Scan (diffScan) where++import Control.Monad+import Futhark.AD.Rev.Monad+import Futhark.Analysis.PrimExp.Convert+import Futhark.Builder+import Futhark.IR.SOACS+import Futhark.Tools+import Futhark.Transform.Rename+import Futhark.Util (pairs, unpairs)++data FirstOrSecond = WrtFirst | WrtSecond++-- computes `d(x op y)/dx` or d(x op y)/dy+mkScanAdjointLam :: VjpOps -> Lambda SOACS -> FirstOrSecond -> ADM (Lambda SOACS)+mkScanAdjointLam ops lam0 which = do+  let len = length $ lambdaReturnType lam0+  lam <- renameLambda lam0+  let p2diff =+        case which of+          WrtFirst -> take len $ lambdaParams lam+          WrtSecond -> drop len $ lambdaParams lam+  p_adjs <- mapM unitAdjOfType (lambdaReturnType lam)+  vjpLambda ops p_adjs (map paramName p2diff) lam++-- Should generate something like:+-- `\ j -> let i = n - 1 - j+--         if i < n-1 then ( ys_adj[i], df2dx ys[i] xs[i+1]) else (0,1) )`+-- where `ys` is  the result of scan+--       `xs` is  the input  of scan+--       `ys_adj` is the known adjoint of ys+--       `j` draw values from `iota n`+mkScanFusedMapLam :: VjpOps -> SubExp -> Lambda SOACS -> [VName] -> [VName] -> [VName] -> ADM (Lambda SOACS)+mkScanFusedMapLam ops w scn_lam xs ys ys_adj = do+  lam <- mkScanAdjointLam ops scn_lam WrtFirst+  ys_ts <- mapM lookupType ys+  par_i <- newParam "i" $ Prim int64+  let i = paramName par_i+  mkLambda [par_i] $+    fmap varsRes . letTupExp "x"+      =<< eIf+        (toExp $ le64 i .==. 0)+        ( buildBody_ $ do+            zs <- mapM (letSubExp "ct_zero" . zeroExp . rowType) ys_ts+            os <- mapM (letSubExp "ct_one" . oneExp . rowType) ys_ts+            pure $ subExpsRes $ unpairs $ zip zs os+        )+        ( buildBody_ $ do+            j <- letSubExp "j" =<< toExp (pe64 w - (le64 i + 1))+            j1 <- letSubExp "j1" =<< toExp (pe64 w - le64 i)+            let index idx arr t = BasicOp $ Index arr $ fullSlice t [DimFix idx]+            y_s <- forM (zip ys_adj ys_ts) $ \(y_, t) ->+              letSubExp (baseString y_ ++ "_j") $ index j y_ t+            lam_rs <-+              eLambda lam . map pure $+                zipWith (index j) ys ys_ts ++ zipWith (index j1) xs ys_ts+            pure $ unpairs $ zip (subExpsRes y_s) lam_rs+        )++-- \(a1, b1) (a2, b2) -> (a2 + b2 * a1, b1 * b2)+mkScanLinFunO :: Type -> ADM (Scan SOACS)+mkScanLinFunO t = do+  let pt = elemType t+  zero <- letSubExp "zeros" $ zeroExp t+  one <- letSubExp "ones" $ oneExp t+  tmp <- mapM newVName ["a1", "b1", "a2", "b2"]+  let [a1, b1, a2, b2] = tmp+      pet = primExpFromSubExp pt . Var+  lam <- mkLambda (map (\v -> Param mempty v t) [a1, b1, a2, b2]) . fmap varsRes $+    tabNest (arrayRank t) [a1, b1, a2, b2] $ \_ [a1', b1', a2', b2'] -> do+      x <- letExp "x" <=< toExp $ pet a2' ~+~ pet b2' ~*~ pet a1'+      y <- letExp "y" <=< toExp $ pet b1' ~*~ pet b2'+      pure [x, y]+  pure $ Scan lam [zero, one]++-- build the map following the scan with linear-function-composition:+-- for each (ds,cs) length-n array results of the scan, combine them as:+--    `let rs = map2 (\ d_i c_i -> d_i + c_i * y_adj[n-1]) d c |> reverse`+-- but insert explicit indexing to reverse inside the map.+mkScan2ndMaps :: SubExp -> (Type, VName, (VName, VName)) -> ADM VName+mkScan2ndMaps w (arr_tp, y_adj, (ds, cs)) = do+  nm1 <- letSubExp "nm1" =<< toExp (pe64 w - 1)+  y_adj_last <-+    letExp (baseString y_adj ++ "_last") $+      BasicOp $ Index y_adj $ fullSlice arr_tp [DimFix nm1]++  par_i <- newParam "i" $ Prim int64+  lam <- mkLambda [par_i] $ do+    let i = paramName par_i+    j <- letSubExp "j" =<< toExp (pe64 w - (le64 i + 1))+    dj <- letExp (baseString ds ++ "_dj") $ BasicOp $ Index ds $ fullSlice arr_tp [DimFix j]+    cj <- letExp (baseString cs ++ "_cj") $ BasicOp $ Index cs $ fullSlice arr_tp [DimFix j]++    let pet = primExpFromSubExp (elemType arr_tp) . Var+    fmap varsRes . tabNest (arrayRank (rowType arr_tp)) [y_adj_last, dj, cj] $ \_ [y_adj_last', dj', cj'] ->+      letTupExp "res" <=< toExp $ pet dj' ~+~ pet cj' ~*~ pet y_adj_last'++  iota <- letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64+  letExp "after_scan" $ Op (Screma w [iota] (ScremaForm [] [] lam))++-- perform the final map, which is fusable with the maps obtained from `mkScan2ndMaps`+-- let xs_contribs =+--    map3 (\ i a r -> if i==0 then r else (df2dy (ys[i-1]) a) \bar{*} r)+--         (iota n) xs rs+mkScanFinalMap :: VjpOps -> SubExp -> Lambda SOACS -> [VName] -> [VName] -> [VName] -> ADM [VName]+mkScanFinalMap ops w scan_lam xs ys rs = do+  let eltps = lambdaReturnType scan_lam+  lam <- mkScanAdjointLam ops scan_lam WrtSecond+  par_i <- newParam "i" $ Prim int64+  let i = paramName par_i+  par_x <- mapM (\(x, t) -> newParam (baseString x ++ "_par_x") t) $ zip xs eltps+  par_r <- mapM (\(r, t) -> newParam (baseString r ++ "_par_r") t) $ zip rs eltps++  map_lam <-+    mkLambda (par_i : par_x ++ par_r) $+      fmap varsRes . letTupExp "scan_contribs"+        =<< eIf+          (toExp $ le64 i .==. 0)+          (resultBodyM $ map (Var . paramName) par_r)+          ( buildBody_ $ do+              im1 <- letSubExp "im1" =<< toExp (le64 i - 1)+              ys_im1 <- forM ys $ \y -> do+                y_t <- lookupType y+                letSubExp (baseString y ++ "_last") $ BasicOp $ Index y $ fullSlice y_t [DimFix im1]++              lam_res <-+                mapM (letExp "const" . BasicOp . SubExp . resSubExp)+                  =<< eLambda lam (map eSubExp $ ys_im1 ++ map (Var . paramName) par_x)++              fmap (varsRes . mconcat) . forM (zip3 lam_res (map paramName par_r) eltps) $+                \(lam_r, r, eltp) -> do+                  let pet = primExpFromSubExp (elemType eltp) . Var++                  tabNest (arrayRank eltp) [lam_r, r] $ \_ [lam_r', r'] ->+                    letTupExp "res" <=< toExp $ pet lam_r' ~*~ pet r'+          )++  iota <- letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64+  letTupExp "scan_contribs" $ Op (Screma w (iota : xs ++ rs) (ScremaForm [] [] map_lam))++diffScan :: VjpOps -> [VName] -> SubExp -> [VName] -> Scan SOACS -> ADM ()+diffScan ops ys w as scan = do+  ys_adj <- mapM lookupAdjVal ys+  as_ts <- mapM lookupType as+  map1_lam <- mkScanFusedMapLam ops w (scanLambda scan) as ys ys_adj+  scans_lin_fun_o <- mapM mkScanLinFunO $ lambdaReturnType $ scanLambda scan+  iota <-+    letExp "iota" $ BasicOp $ Iota w (intConst Int64 0) (intConst Int64 1) Int64+  r_scan <-+    letTupExp "adj_ctrb_scan" $+      Op (Screma w [iota] (ScremaForm scans_lin_fun_o [] map1_lam))+  red_nms <- mapM (mkScan2ndMaps w) (zip3 as_ts ys_adj (pairs r_scan))+  as_contribs <- mkScanFinalMap ops w (scanLambda scan) as ys red_nms+  zipWithM_ updateAdj as as_contribs

@@ -0,0 +1,188 @@+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeFamilies #-}++module Futhark.AD.Rev.Scatter (vjpScatter) where++import Control.Monad+import Futhark.AD.Rev.Monad+import Futhark.Analysis.PrimExp.Convert+import Futhark.Builder+import Futhark.IR.SOACS+import Futhark.Tools+import Futhark.Util (chunk)++withinBounds :: [(SubExp, VName)] -> TPrimExp Bool VName+withinBounds [] = TPrimExp $ ValueExp (BoolValue True)+withinBounds [(q, i)] = (le64 i .<. pe64 q) .&&. (pe64 (intConst Int64 (-1)) .<. le64 i)+withinBounds (qi : qis) = withinBounds [qi] .&&. withinBounds qis++-- Generates a potential tower-of-maps lambda body for an indexing operation.+-- Assuming parameters:+--   `arr`   the array that is indexed+--   `[(w_1, i_1), (w_2, i_2), ..., (w_k, i_k)]` outer lambda formal parameters and their bounds+--   `[n_1,n_2,...]ptp` the type of the index expression `arr[i_1,i_2,...,i_k]`+-- Generates something like:+-- (\ i_1 i_2 ->+--    map (\j_1 -> ... if (i_1 >= 0 && i_1 < w_1) &&+--                        (i_2 >= 0 && i_2 < w_2) && ...+--                     then arr[i_1, i_2, ... j_1, ...]+--                     else 0+--        ) (iota n_1)+-- )+-- The idea is that you do not want to put under the `if` something+--     that is an array because it would not flatten well!+genIdxLamBody :: VName -> [(SubExp, Param Type)] -> Type -> ADM (Body SOACS)+genIdxLamBody as wpis = genRecLamBody as wpis []+  where+    genRecLamBody :: VName -> [(SubExp, Param Type)] -> [Param Type] -> Type -> ADM (Body SOACS)+    genRecLamBody arr w_pis nest_pis (Array t (Shape []) _) =+      genRecLamBody arr w_pis nest_pis (Prim t)+    genRecLamBody arr w_pis nest_pis (Array t (Shape (s : ss)) _) = do+      new_ip <- newParam "i" (Prim int64)+      let t' = Prim t `arrayOfShape` Shape ss+      inner_lam <-+        mkLambda [new_ip] $+          bodyBind =<< genRecLamBody arr w_pis (nest_pis ++ [new_ip]) t'+      let (_, orig_pis) = unzip w_pis+      buildBody_ . localScope (scopeOfLParams (orig_pis ++ nest_pis)) $ do+        iota_v <- letExp "iota" $ BasicOp $ Iota s (intConst Int64 0) (intConst Int64 1) Int64+        r <- letSubExp (baseString arr ++ "_elem") $ Op $ Screma s [iota_v] (mapSOAC inner_lam)+        pure [subExpRes r]+    genRecLamBody arr w_pis nest_pis (Prim ptp) = do+      let (ws, orig_pis) = unzip w_pis+      let inds = map paramName (orig_pis ++ nest_pis)+      localScope (scopeOfLParams (orig_pis ++ nest_pis)) $+        eBody+          [ eIf+              (toExp $ withinBounds $ zip ws $ map paramName orig_pis)+              ( do+                  r <- letSubExp "r" $ BasicOp $ Index arr $ Slice $ map (DimFix . Var) inds+                  resultBodyM [r]+              )+              (resultBodyM [Constant $ blankPrimValue ptp])+          ]+    genRecLamBody _ _ _ _ = error "In Rev.hs, helper function genRecLamBody, unreachable case reached!"++--+-- Original:+--   let ys = scatter xs is vs+-- Assumes no duplicate indices in `is`+-- Forward Sweep:+--   let xs_save = gather xs is+--   let ys = scatter xs is vs+-- Return Sweep:+--   let vs_ctrbs = gather is ys_adj+--   let vs_adj \overline{+}= vs_ctrbs -- by map or generalized reduction+--   let xs_adj = scatter ys_adj is \overline{0}+--   let xs = scatter ys is xs_save+vjpScatter1 ::+  PatElem Type ->+  StmAux () ->+  (SubExp, [VName], (ShapeBase SubExp, Int, VName)) ->+  ADM () ->+  ADM ()+vjpScatter1 pys aux (w, ass, (shp, num_vals, xs)) m = do+  let rank = length $ shapeDims shp+      (all_inds, val_as) = splitAt (rank * num_vals) ass+      inds_as = chunk rank all_inds+  xs_t <- lookupType xs+  let val_t = stripArray (shapeRank shp) xs_t+  -- computing xs_save+  xs_saves <- mkGather inds_as xs xs_t+  -- performing the scatter+  id_lam <-+    mkIdentityLambda $+      replicate (shapeRank shp) (Prim int64) ++ replicate (shapeRank shp) val_t+  addStm $ Let (Pat [pys]) aux $ Op $ Scatter w ass id_lam [(shp, num_vals, xs)]+  m+  let ys = patElemName pys+  -- XXX: Since our restoration of xs will consume ys, we have to+  -- make a copy of ys in the chance that it is actually the result+  -- of the program.  In that case the asymptotics will not be+  -- (locally) preserved, but since ys must necessarily have been+  -- constructed somewhere close, they are probably globally OK.+  ys_copy <- letExp (baseString ys <> "_copy") $ BasicOp $ Copy ys+  returnSweepCode $ do+    ys_adj <- lookupAdjVal ys+    -- computing vs_ctrbs and updating vs_adj+    vs_ctrbs <- mkGather inds_as ys_adj xs_t+    zipWithM_ updateAdj val_as vs_ctrbs -- use Slice?+    -- creating xs_adj+    zeros <-+      replicateM (length val_as) . letExp "zeros" $+        zeroExp $ xs_t `setOuterSize` w+    let f_tps = replicate (rank * num_vals) (Prim int64) ++ replicate num_vals val_t+    f <- mkIdentityLambda f_tps+    xs_adj <-+      letExp (baseString xs ++ "_adj") . Op $+        Scatter w (all_inds ++ zeros) f [(shp, num_vals, ys_adj)]+    insAdj xs xs_adj -- reusing the ys_adj for xs_adj!+    f' <- mkIdentityLambda f_tps+    xs_rc <-+      auxing aux . letExp (baseString xs <> "_rc") . Op $+        Scatter w (all_inds ++ xs_saves) f' [(shp, num_vals, ys)]+    addSubstitution xs xs_rc+    addSubstitution ys ys_copy+  where+    -- Creates a potential map-nest that indexes in full the array,+    --   and applies the condition of indices within bounds at the+    --   deepest level in the nest so that everything can be parallel.+    mkGather :: [[VName]] -> VName -> Type -> ADM [VName]+    mkGather inds_as arr arr_t = do+      ips <- forM inds_as $ \idxs ->+        mapM (\idx -> newParam (baseString idx ++ "_elem") (Prim int64)) idxs++      gather_lam <- mkLambda (concat ips) . fmap mconcat . forM ips $ \idxs -> do+        let q = length idxs+            (ws, eltp) = (take q $ arrayDims arr_t, stripArray q arr_t)+        bodyBind =<< genIdxLamBody arr (zip ws idxs) eltp+      let soac = Screma w (concat inds_as) (mapSOAC gather_lam)+      letTupExp (baseString arr ++ "_gather") $ Op soac++vjpScatter ::+  VjpOps ->+  Pat Type ->+  StmAux () ->+  (SubExp, [VName], Lambda SOACS, [(Shape, Int, VName)]) ->+  ADM () ->+  ADM ()+vjpScatter ops (Pat pes) aux (w, ass, lam, written_info) m+  | isIdentityLambda lam,+    [(shp, num_vals, xs)] <- written_info,+    [pys] <- pes =+      vjpScatter1 pys aux (w, ass, (shp, num_vals, xs)) m+  | isIdentityLambda lam = do+      let sind = splitInd written_info+          (inds, vals) = splitAt sind ass+      lst_stms <- chunkScatterInps (inds, vals) (zip pes written_info)+      diffScatters (stmsFromList lst_stms)+  | otherwise =+      error "vjpScatter: cannot handle"+  where+    splitInd [] = 0+    splitInd ((shp, num_res, _) : rest) =+      num_res * length (shapeDims shp) + splitInd rest+    chunkScatterInps (acc_inds, acc_vals) [] =+      case (acc_inds, acc_vals) of+        ([], []) -> pure []+        _ -> error "chunkScatterInps: cannot handle"+    chunkScatterInps+      (acc_inds, acc_vals)+      ((pe, info@(shp, num_vals, _)) : rest) = do+        let num_inds = num_vals * length (shapeDims shp)+            (curr_inds, other_inds) = splitAt num_inds acc_inds+            (curr_vals, other_vals) = splitAt num_vals acc_vals+        vtps <- mapM lookupType curr_vals+        f <- mkIdentityLambda (replicate num_inds (Prim int64) ++ vtps)+        let stm =+              Let (Pat [pe]) aux . Op $+                Scatter w (curr_inds ++ curr_vals) f [info]+        stms_rest <- chunkScatterInps (other_inds, other_vals) rest+        pure $ stm : stms_rest+    diffScatters all_stms+      | Just (stm, stms) <- stmsHead all_stms =+          vjpStm ops stm $ diffScatters stms+      | otherwise = m

@@ -67,6 +67,10 @@   (lumap'', used'') <- analyseKernelBody (lumap', used') body   let nms = (freeIn lvl <> freeIn tps) `namesSubtract` used''   pure (insertNames pat_name nms lumap'', used'' <> nms)+analyseGPUOp pat_name (lumap, used) (Inner (GPUBody ts body)) = do+  (lumap', used') <- analyseBody lumap used body+  let nms = freeIn ts+  pure (insertNames pat_name nms lumap', used' <> nms)  segOpHelper ::   (FreeIn (OpWithAliases (Op rep)), ASTRep rep) =>

@@ -16,6 +16,7 @@     MetricsM,     stmMetrics,     lambdaMetrics,+    bodyMetrics,   ) where @@ -91,6 +92,7 @@ funDefMetrics :: OpMetrics (Op rep) => FunDef rep -> MetricsM () funDefMetrics = bodyMetrics . funDefBody +-- | Compute metrics for this body. bodyMetrics :: OpMetrics (Op rep) => Body rep -> MetricsM () bodyMetrics = mapM_ stmMetrics . bodyStms 

@@ -55,8 +55,18 @@     zExt32,     zExt64,     sExtAs,+    fMin16,+    fMin32,     fMin64,+    fMax16,+    fMax32,     fMax64,++    -- * Untyped construction+    (~*~),+    (~/~),+    (~+~),+    (~-~),   ) where @@ -429,6 +439,15 @@         TPrimExp $ constFoldPrimExp z     | otherwise = numBad "divRoundingUp" (x, y) +  TPrimExp x `pow` TPrimExp y+    | Just z <-+        msum+          [ asIntOp Pow x y,+            asFloatOp FPow x y+          ] =+        TPrimExp $ constFoldPrimExp z+    | otherwise = numBad "pow" (x, y)+   sgn (TPrimExp (ValueExp (IntValue i))) = Just $ signum $ valueIntegral i   sgn _ = Nothing @@ -628,13 +647,29 @@ zExt64 :: IntExp t => TPrimExp t v -> TPrimExp Int64 v zExt64 = isInt64 . zExt Int64 . untyped +-- | 16-bit float minimum.+fMin16 :: TPrimExp Half v -> TPrimExp Half v -> TPrimExp Half v+fMin16 x y = isF16 $ BinOpExp (FMin Float16) (untyped x) (untyped y)++-- | 32-bit float minimum.+fMin32 :: TPrimExp Float v -> TPrimExp Float v -> TPrimExp Float v+fMin32 x y = isF32 $ BinOpExp (FMin Float32) (untyped x) (untyped y)+ -- | 64-bit float minimum. fMin64 :: TPrimExp Double v -> TPrimExp Double v -> TPrimExp Double v-fMin64 x y = TPrimExp $ BinOpExp (FMin Float64) (untyped x) (untyped y)+fMin64 x y = isF64 $ BinOpExp (FMin Float64) (untyped x) (untyped y) +-- | 16-bit float maximum.+fMax16 :: TPrimExp Half v -> TPrimExp Half v -> TPrimExp Half v+fMax16 x y = isF16 $ BinOpExp (FMax Float16) (untyped x) (untyped y)++-- | 32-bit float maximum.+fMax32 :: TPrimExp Float v -> TPrimExp Float v -> TPrimExp Float v+fMax32 x y = isF32 $ BinOpExp (FMax Float32) (untyped x) (untyped y)+ -- | 64-bit float maximum. fMax64 :: TPrimExp Double v -> TPrimExp Double v -> TPrimExp Double v-fMax64 x y = TPrimExp $ BinOpExp (FMax Float64) (untyped x) (untyped y)+fMax64 x y = isF64 $ BinOpExp (FMax Float64) (untyped x) (untyped y)  -- | Convert result of some integer expression to have the same type -- as another, using sign extension.@@ -672,3 +707,53 @@   S.union (leafExpTypes e1) (leafExpTypes e2) leafExpTypes (FunExp _ pes _) =   S.unions $ map leafExpTypes pes++-- | Multiplication of untyped 'PrimExp's, which must have the same+-- type.+(~*~) :: PrimExp v -> PrimExp v -> PrimExp v+x ~*~ y = BinOpExp op x y+  where+    t = primExpType x+    op = case t of+      IntType it -> Mul it OverflowUndef+      FloatType ft -> FMul ft+      Bool -> LogAnd+      Unit -> LogAnd++-- | Division of untyped 'PrimExp's, which must have the same+-- type.  For integers, this is unsafe signed division.+(~/~) :: PrimExp v -> PrimExp v -> PrimExp v+x ~/~ y = BinOpExp op x y+  where+    t = primExpType x+    op = case t of+      IntType it -> SDiv it Unsafe+      FloatType ft -> FDiv ft+      Bool -> LogAnd+      Unit -> LogAnd++-- | Addition of untyped 'PrimExp's, which must have the same type.+(~+~) :: PrimExp v -> PrimExp v -> PrimExp v+x ~+~ y = BinOpExp op x y+  where+    t = primExpType x+    op = case t of+      IntType it -> Add it OverflowUndef+      FloatType ft -> FAdd ft+      Bool -> LogOr+      Unit -> LogOr++-- | Subtraction of untyped 'PrimExp's, which must have the same type.+(~-~) :: PrimExp v -> PrimExp v -> PrimExp v+x ~-~ y = BinOpExp op x y+  where+    t = primExpType x+    op = case t of+      IntType it -> Sub it OverflowUndef+      FloatType ft -> FSub ft+      Bool -> LogOr+      Unit -> LogOr++infix 7 ~*~, ~/~++infix 6 ~+~, ~-~

@@ -141,7 +141,8 @@  -- | How to run a benchmark. data RunOptions = RunOptions-  { runMinRuns :: Int,+  { -- | Applies both to initial and convergence phase.+    runMinRuns :: Int,     runMinTime :: NominalDiffTime,     runTimeout :: Int,     runVerbose :: Int,@@ -154,8 +155,8 @@     runResultAction :: (Int, Maybe Double) -> IO ()   } --- | A list of @(autocorrelation,rsd)@ pairs.  When the--- autocorrelation is above the first element and the RSD is above the+-- | A list of @(autocorrelation,rse)@ pairs.  When the+-- autocorrelation is above the first element and the RSE is above the -- second element, we want more runs. convergenceCriteria :: [(Double, Double)] convergenceCriteria =@@ -168,9 +169,9 @@  -- Returns the next run count. nextRunCount :: Int -> Double -> Double -> Int-nextRunCount runs rsd acor = if any check convergenceCriteria then div runs 2 else 0+nextRunCount runs rse acor = if any check convergenceCriteria then div runs 2 else 0   where-    check (acor_lb, rsd_lb) = acor > acor_lb && rsd > rsd_lb+    check (acor_lb, rse_lb) = acor > acor_lb && rse > rse_lb  type BenchM = ExceptT T.Text IO @@ -216,7 +217,7 @@   BenchM (DL.DList (RunResult, [T.Text])) runConvergence do_run opts initial_r =   let runtimes = resultRuntimes (DL.toList initial_r)-      (n, _, rsd, acor) = runtimesMetrics runtimes+      (n, _, rse, acor) = runtimesMetrics runtimes    in -- If the runtimes collected during the runMinimum phase are       -- unstable enough that we need more in order to converge, we throw       -- away the runMinimum runtimes.  This is because they often exhibit@@ -224,11 +225,11 @@       -- outliers that poison the metrics we use to determine convergence.       -- By throwing them away we converge much faster, and still get the       -- right result.-      case nextRunCount n rsd acor of+      case nextRunCount n rse acor of         x           | x > 0,             runConvergencePhase opts ->-              moreRuns mempty mempty rsd x+              moreRuns mempty mempty rse (x `max` runMinRuns opts)           | otherwise ->               pure initial_r   where@@ -237,30 +238,30 @@      runtimesMetrics runtimes =       let n = U.length runtimes-          rsd = (fastStdDev runtimes / sqrt (fromIntegral n)) / mean runtimes+          rse = (fastStdDev runtimes / sqrt (fromIntegral n)) / mean runtimes           (x, _, _) = autocorrelation runtimes        in ( n,             realToFrac (U.sum runtimes) :: NominalDiffTime,-            rsd,+            rse,             fromMaybe 1 (x U.!? 1)           ) -    sample rsd = do+    sample rse = do       x <- do_run-      liftIO $ runResultAction opts (runMicroseconds (fst x), Just rsd)+      liftIO $ runResultAction opts (runMicroseconds (fst x), Just rse)       pure x -    moreRuns runtimes r rsd x = do-      r' <- replicateM x $ sample rsd+    moreRuns runtimes r rse x = do+      r' <- replicateM x $ sample rse       loop (runtimes <> resultRuntimes r') (r <> DL.fromList r')      loop runtimes r = do-      let (n, total, rsd, acor) = runtimesMetrics runtimes-      case nextRunCount n rsd acor of+      let (n, total, rse, acor) = runtimesMetrics runtimes+      case nextRunCount n rse acor of         x           | x > 0,             total < runConvergenceMaxTime opts ->-              moreRuns runtimes r rsd x+              moreRuns runtimes r rse x           | otherwise ->               pure r 

@@ -18,14 +18,17 @@ import Data.Maybe import Data.Ord import qualified Data.Text as T+import qualified Data.Text.IO as T import Data.Time.Clock (NominalDiffTime, UTCTime, diffUTCTime, getCurrentTime) import qualified Data.Vector.Unboxed as U import Futhark.Bench import Futhark.Server import Futhark.Test-import Futhark.Util (atMostChars, fancyTerminal, maybeNth, pmapIO)+import Futhark.Util (atMostChars, fancyTerminal, pmapIO) import Futhark.Util.Console import Futhark.Util.Options+import Futhark.Util.Pretty (prettyText)+import Futhark.Util.ProgressBar import Statistics.Resampling (Estimator (..), resample) import Statistics.Resampling.Bootstrap (bootstrapBCA) import Statistics.Types (cl95, confIntLDX, confIntUDX, estError, estPoint)@@ -248,50 +251,36 @@       runResultAction = f     } -progressBar :: Double -> Double -> Int -> String-progressBar cur bound steps =-  "|" <> map cell [1 .. steps] <> "| "-  where-    step_size :: Double-    step_size = bound / fromIntegral steps-    chars = " ▏▎▍▍▌▋▊▉█"-    char i = fromMaybe ' ' $ maybeNth (i :: Int) chars-    num_chars = fromIntegral $ length chars--    cell :: Int -> Char-    cell i-      | i' * step_size <= cur = char 9-      | otherwise =-          char (floor (((cur - (i' - 1) * step_size) * num_chars) / step_size))-      where-        i' = fromIntegral i- descString :: String -> Int -> String descString desc pad_to = desc ++ ": " ++ replicate (pad_to - length desc) ' ' -progressBarSteps :: Int-progressBarSteps = 10+progress :: Double -> T.Text+progress elapsed =+  progressBar+    ( ProgressBar+        { progressBarSteps = 10,+          progressBarBound = 1,+          progressBarElapsed = elapsed+        }+    ) -interimResult :: Int -> Int -> Double -> Double -> String-interimResult us_sum runs elapsed bound =-  printf "%10.0fμs " avg-    <> progressBar elapsed bound progressBarSteps-    <> (" " <> show runs <> " runs")+interimResult :: Int -> Int -> Double -> T.Text+interimResult us_sum runs elapsed =+  T.pack (printf "%10.0fμs " avg) <> progress elapsed+    <> (" " <> prettyText runs <> " runs")   where     avg :: Double     avg = fromIntegral us_sum / fromIntegral runs -convergenceBar :: (String -> IO ()) -> IORef Int -> Int -> Int -> Double -> IO ()-convergenceBar p spin_count us_sum i rsd' = do+convergenceBar :: (T.Text -> IO ()) -> IORef Int -> Int -> Int -> Double -> IO ()+convergenceBar p spin_count us_sum i rse' = do   spin_idx <- readIORef spin_count-  let spin_char = spin_load !! spin_idx-  p $ printf "%10.0fμs %c (RSD of mean: %2.4f; %4d runs)" avg spin_char rsd' i-  let spin_count' = (spin_idx + 1) `mod` 10-  writeIORef spin_count spin_count'+  let spin = progressSpinner spin_idx+  p $ T.pack $ printf "%10.0fμs %s (RSE of mean: %2.4f; %4d runs)" avg spin rse' i+  writeIORef spin_count (spin_idx + 1)   where     avg :: Double     avg = fromIntegral us_sum / fromIntegral i-    spin_load = "⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏"  data BenchPhase = Initial | Convergence @@ -301,11 +290,11 @@       count <- newIORef (0, 0)       phase_var <- newIORef Initial       spin_count <- newIORef 0-      pure $ \(us, rsd) -> do+      pure $ \(us, rse) -> do         putStr "\r" -- Go to start of line.         let p s =-              putStr $-                descString (atMostChars maxDatasetNameLength dataset_desc) pad_to ++ s+              T.putStr $+                T.pack (descString (atMostChars maxDatasetNameLength dataset_desc) pad_to) <> s          (us_sum, i) <- readIORef count @@ -322,27 +311,27 @@          phase <- readIORef phase_var -        case (us, phase, rsd) of+        case (us, phase, rse) of           (Nothing, _, _) ->             let elapsed = determineProgress i-             in p $ replicate 13 ' ' <> progressBar elapsed 1.0 progressBarSteps+             in p $ T.pack (replicate 13 ' ') <> progress elapsed           (Just us', Initial, Nothing) -> do             let us_sum' = us_sum + us'                 i' = i + 1             writeIORef count (us_sum', i')             let elapsed = determineProgress i'-            p $ interimResult us_sum' i' elapsed 1.0-          (Just us', Initial, Just rsd') -> do+            p $ interimResult us_sum' i' elapsed+          (Just us', Initial, Just rse') -> do             -- Switched from phase 1 to convergence; discard all             -- prior results.             writeIORef count (us', 1)             writeIORef phase_var Convergence-            convergenceBar p spin_count us' 1 rsd'-          (Just us', Convergence, Just rsd') -> do+            convergenceBar p spin_count us' 1 rse'+          (Just us', Convergence, Just rse') -> do             let us_sum' = us_sum + us'                 i' = i + 1             writeIORef count (us_sum', i')-            convergenceBar p spin_count us_sum' i' rsd'+            convergenceBar p spin_count us_sum' i' rse'           (Just _, Convergence, Nothing) ->             pure () -- Probably should not happen.         putStr " " -- Just to move the cursor away from the progress bar.

@@ -33,6 +33,7 @@ import Futhark.Optimise.CSE import Futhark.Optimise.DoubleBuffer import Futhark.Optimise.Fusion+import Futhark.Optimise.HistAccs import Futhark.Optimise.InPlaceLowering import Futhark.Optimise.InliningDeadFun import qualified Futhark.Optimise.MemoryBlockMerging as MemoryBlockMerging@@ -40,6 +41,7 @@ import Futhark.Optimise.TileLoops import Futhark.Optimise.Unstream import Futhark.Pass+import Futhark.Pass.AD import Futhark.Pass.ExpandAllocations import qualified Futhark.Pass.ExplicitAllocations.GPU as GPU import qualified Futhark.Pass.ExplicitAllocations.Seq as Seq@@ -549,8 +551,11 @@     soacsPassOption inlineAggressively [],     soacsPassOption inlineConservatively [],     soacsPassOption removeDeadFunctions [],+    soacsPassOption applyAD [],+    soacsPassOption applyADInnermost [],     kernelsPassOption babysitKernels [],     kernelsPassOption tileLoops [],+    kernelsPassOption histAccsGPU [],     kernelsPassOption unstreamGPU [],     kernelsPassOption sinkGPU [],     typedPassOption soacsProg GPU extractKernels [],

@@ -5,8 +5,8 @@ -- | @futhark lsp@ module Futhark.CLI.LSP (main) where -import Control.Concurrent.MVar (newMVar) import Control.Monad.IO.Class (MonadIO (liftIO))+import Data.IORef (newIORef) import Futhark.LSP.Handlers (handlers) import Futhark.LSP.State (emptyState) import Futhark.Util (debug)@@ -22,7 +22,7 @@ -- | Run @futhark lsp@ main :: String -> [String] -> IO () main _prog _args = do-  state_mvar <- newMVar emptyState+  state_mvar <- newIORef emptyState   debug "Init with emptyState"   setupLogger Nothing ["futhark"] DEBUG   _ <-@@ -30,7 +30,7 @@       ServerDefinition         { onConfigurationChange = const $ const $ Right (),           defaultConfig = (),-          doInitialize = \env _req -> do pure $ Right env,+          doInitialize = \env _req -> pure $ Right env,           staticHandlers = handlers state_mvar,           interpretHandler = \env -> Iso (runLspT env) liftIO,           options =

@@ -34,6 +34,7 @@ import Futhark.Test.Values import Futhark.Util   ( directoryContents,+    fancyTerminal,     hashText,     nubOrd,     runProgramWithExitCode,@@ -41,6 +42,8 @@ import Futhark.Util.Options import Futhark.Util.Pretty (prettyText, prettyTextOneLine) import qualified Futhark.Util.Pretty as PP+import Futhark.Util.ProgressBar+import System.Console.ANSI (clearLine) import System.Directory   ( copyFile,     createDirectoryIfMissing,@@ -750,10 +753,33 @@     format = fromMaybe "webm" $ videoFormat params     bmpfile dir j = dir </> printf "frame%010d.bmp" (j :: Int) -    renderFrames dir (stepfun, closure) initial num_frames =+    (progressStep, progressDone)+      | fancyTerminal,+        scriptVerbose (envOpts env) > 0 =+          ( \j num_frames -> do+              liftIO . T.putStr $+                "\r"+                  <> progressBar+                    (ProgressBar 40 (fromIntegral num_frames) (fromIntegral j))+                  <> "generating frame "+                  <> prettyText j+                  <> "/"+                  <> prettyText num_frames+                  <> " "+              liftIO $ hFlush stdout,+            liftIO $ do+              T.putStr "\r"+              clearLine+          )+      | otherwise =+          (\_ _ -> pure (), pure ())++    renderFrames dir (stepfun, closure) initial num_frames = do       foldM_ frame initial [0 .. num_frames - 1]+      progressDone       where         frame old_state j = do+          progressStep j num_frames           v <-             evalExp literateBuiltin (envServer env)               . Call (FuncFut stepfun)

@@ -0,0 +1,144 @@+{-# LANGUAGE OverloadedStrings #-}++-- | The main function for the @futhark@ command line program.+module Futhark.CLI.Main (main) where++import Control.Exception+import Data.List (sortOn)+import Data.Maybe+import qualified Data.Text as T+import qualified Data.Text.IO as T+import qualified Futhark.CLI.Autotune as Autotune+import qualified Futhark.CLI.Bench as Bench+import qualified Futhark.CLI.C as C+import qualified Futhark.CLI.CUDA as CCUDA+import qualified Futhark.CLI.Check as Check+import qualified Futhark.CLI.Datacmp as Datacmp+import qualified Futhark.CLI.Dataset as Dataset+import qualified Futhark.CLI.Defs as Defs+import qualified Futhark.CLI.Dev as Dev+import qualified Futhark.CLI.Doc as Doc+import qualified Futhark.CLI.LSP as LSP+import qualified Futhark.CLI.Literate as Literate+import qualified Futhark.CLI.Misc as Misc+import qualified Futhark.CLI.Multicore as Multicore+import qualified Futhark.CLI.MulticoreWASM as MulticoreWASM+import qualified Futhark.CLI.OpenCL as OpenCL+import qualified Futhark.CLI.Pkg as Pkg+import qualified Futhark.CLI.PyOpenCL as PyOpenCL+import qualified Futhark.CLI.Python as Python+import qualified Futhark.CLI.Query as Query+import qualified Futhark.CLI.REPL as REPL+import qualified Futhark.CLI.Run as Run+import qualified Futhark.CLI.Test as Test+import qualified Futhark.CLI.WASM as WASM+import Futhark.Error+import Futhark.Util (maxinum)+import Futhark.Util.Options+import GHC.IO.Encoding (setLocaleEncoding)+import GHC.IO.Exception (IOErrorType (..), IOException (..))+import System.Environment+import System.Exit+import System.IO+import Prelude++type Command = String -> [String] -> IO ()++commands :: [(String, (Command, String))]+commands =+  sortOn+    fst+    [ ("dev", (Dev.main, "Run compiler passes directly.")),+      ("repl", (REPL.main, "Run interactive Read-Eval-Print-Loop.")),+      ("run", (Run.main, "Run a program through the (slow!) interpreter.")),+      ("c", (C.main, "Compile to sequential C.")),+      ("opencl", (OpenCL.main, "Compile to C calling OpenCL.")),+      ("cuda", (CCUDA.main, "Compile to C calling CUDA.")),+      ("multicore", (Multicore.main, "Compile to multicore C.")),+      ("python", (Python.main, "Compile to sequential Python.")),+      ("pyopencl", (PyOpenCL.main, "Compile to Python calling PyOpenCL.")),+      ("wasm", (WASM.main, "Compile to WASM with sequential C")),+      ("wasm-multicore", (MulticoreWASM.main, "Compile to WASM with multicore C")),+      ("test", (Test.main, "Test Futhark programs.")),+      ("bench", (Bench.main, "Benchmark Futhark programs.")),+      ("dataset", (Dataset.main, "Generate random test data.")),+      ("datacmp", (Datacmp.main, "Compare Futhark data files for equality.")),+      ("dataget", (Misc.mainDataget, "Extract test data.")),+      ("doc", (Doc.main, "Generate documentation for Futhark code.")),+      ("pkg", (Pkg.main, "Manage local packages.")),+      ("check", (Check.main, "Type-check a program.")),+      ("check-syntax", (Misc.mainCheckSyntax, "Syntax-check a program.")),+      ("imports", (Misc.mainImports, "Print all non-builtin imported Futhark files.")),+      ("hash", (Misc.mainHash, "Print hash of program AST.")),+      ("autotune", (Autotune.main, "Autotune threshold parameters.")),+      ("defs", (Defs.main, "Show location and name of all definitions.")),+      ("query", (Query.main, "Query semantic information about program.")),+      ("literate", (Literate.main, "Process a literate Futhark program.")),+      ("lsp", (LSP.main, "Run LSP server.")),+      ("thanks", (Misc.mainThanks, "Express gratitude."))+    ]++msg :: String+msg =+  unlines $+    ["<command> options...", "Commands:", ""]+      ++ [ "   " <> cmd <> replicate (k - length cmd) ' ' <> desc+           | (cmd, (_, desc)) <- commands+         ]+  where+    k = maxinum (map (length . fst) commands) + 3++-- | Catch all IO exceptions and print a better error message if they+-- happen.+reportingIOErrors :: IO () -> IO ()+reportingIOErrors =+  flip+    catches+    [ Handler onExit,+      Handler onICE,+      Handler onIOException,+      Handler onError+    ]+  where+    onExit :: ExitCode -> IO ()+    onExit = throwIO++    onICE :: InternalError -> IO ()+    onICE (Error CompilerLimitation s) = do+      T.hPutStrLn stderr "Known compiler limitation encountered.  Sorry."+      T.hPutStrLn stderr "Revise your program or try a different Futhark compiler."+      T.hPutStrLn stderr s+      exitWith $ ExitFailure 1+    onICE (Error CompilerBug s) = do+      T.hPutStrLn stderr "Internal compiler error."+      T.hPutStrLn stderr "Please report this at https://github.com/diku-dk/futhark/issues."+      T.hPutStrLn stderr s+      exitWith $ ExitFailure 1++    onError :: SomeException -> IO ()+    onError e+      | Just UserInterrupt <- asyncExceptionFromException e =+          pure () -- This corresponds to CTRL-C, which is not an error.+      | otherwise = do+          T.hPutStrLn stderr "Internal compiler error (unhandled IO exception)."+          T.hPutStrLn stderr "Please report this at https://github.com/diku-dk/futhark/issues"+          T.hPutStrLn stderr $ T.pack $ show e+          exitWith $ ExitFailure 1++    onIOException :: IOException -> IO ()+    onIOException e+      | ioe_type e == ResourceVanished =+          exitWith $ ExitFailure 1+      | otherwise = throw e++main :: IO ()+main = reportingIOErrors $ do+  hSetEncoding stdout utf8+  hSetEncoding stderr utf8+  setLocaleEncoding utf8+  args <- getArgs+  prog <- getProgName+  case args of+    cmd : args'+      | Just (m, _) <- lookup cmd commands -> m (unwords [prog, cmd]) args'+    _ -> mainWithOptions () [] msg (const . const Nothing) prog args

@@ -99,6 +99,17 @@   context prog_ctx $     pureTestResults $ liftIO $ withServer (futharkServerCfg to_run to_run_args) f +data TestMode+  = -- | Only type check.+    TypeCheck+  | -- | Only compile (do not run).+    Compile+  | -- | Test compiled code.+    Compiled+  | -- | Test interpreted code.+    Interpreted+  deriving (Eq, Show)+ data TestCase = TestCase   { _testCaseMode :: TestMode,     testCaseProgram :: FilePath,@@ -242,13 +253,13 @@       let backend = configBackend progs           extra_compiler_options = configExtraCompilerOptions progs -      unless (mode == Compile) $+      when (mode `elem` [Compiled, Interpreted]) $         context "Generating reference outputs" $           -- We probably get the concurrency at the test program level,           -- so force just one data set at a time here.           ensureReferenceOutput (Just 1) (FutharkExe futhark) "c" program ios -      unless (mode == Interpreted) $+      when (mode == Compiled) $         context ("Compiling with --backend=" <> T.pack backend) $ do           compileTestProgram extra_compiler_options (FutharkExe futhark) backend program warnings           mapM_ (testMetrics progs program) structures@@ -265,7 +276,7 @@                 let run = runCompiledEntry (FutharkExe futhark) server program                 concat <$> mapM run ios -      unless (mode == Compile || mode == Compiled) $+      when (mode == Interpreted) $         context "Interpreting" $           accErrors_ $ map (runInterpretedEntry (FutharkExe futhark) program) ios @@ -323,8 +334,8 @@   | not (match regex $ T.unpack err) =       E.throwError $         "Expected error:\n  " <> regex_s-          <> "\nGot error:\n  "-          <> err+          <> "\nGot error:\n"+          <> T.unlines (map ("  " <>) (T.lines err)) checkError _ _ =   pure () @@ -607,7 +618,7 @@ defaultConfig :: TestConfig defaultConfig =   TestConfig-    { configTestMode = Everything,+    { configTestMode = Compiled,       configExclude = ["disable"],       configPrograms =         ProgConfig@@ -657,14 +668,6 @@ addOption :: String -> ProgConfig -> ProgConfig addOption option config =   config {configExtraOptions = configExtraOptions config ++ [option]}--data TestMode-  = TypeCheck-  | Compile-  | Compiled-  | Interpreted-  | Everything-  deriving (Eq, Show)  commandLineOptions :: [FunOptDescr TestConfig] commandLineOptions =

@@ -29,6 +29,7 @@   ) import Futhark.MonadFreshNames import qualified Language.C.Quote.OpenCL as C+import qualified Language.C.Syntax as C import NeatInterpolation (untrimming)  -- | Compile the program to C with calls to CUDA.@@ -137,7 +138,26 @@            }        ] +-- We detect the special case of writing a constant and turn it into a+-- non-blocking write.  This may be slightly faster, as it prevents+-- unnecessary synchronisation of the context, and writing a constant+-- is fairly common.  This is only possible because we can give the+-- constant infinite lifetime (with 'static'), which is not the case+-- for ordinary variables. writeCUDAScalar :: GC.WriteScalar OpenCL ()+writeCUDAScalar mem idx t "device" _ val@C.Const {} = do+  val' <- newVName "write_static"+  let (bef, aft) = profilingEnclosure copyScalarToDev+  GC.item+    [C.citem|{static $ty:t $id:val' = $exp:val;+              $items:bef+              CUDA_SUCCEED_OR_RETURN(+                cuMemcpyHtoDAsync($exp:mem + $exp:idx * sizeof($ty:t),+                                  &$id:val',+                                  sizeof($ty:t),+                                  0));+              $items:aft+             }|] writeCUDAScalar mem idx t "device" _ val = do   val' <- newVName "write_tmp"   let (bef, aft) = profilingEnclosure copyScalarToDev@@ -167,10 +187,12 @@           cuMemcpyDtoH(&$id:val,                        $exp:mem + $exp:idx * sizeof($ty:t),                        sizeof($ty:t)));-        $items:aft+       $items:aft        }        |]-  GC.stm [C.cstm|if (futhark_context_sync(ctx) != 0) { return 1; }|]+  GC.stm+    [C.cstm|if (ctx->failure_is_an_option && futhark_context_sync(ctx) != 0)+            { return 1; }|]   pure [C.cexp|$id:val|] readCUDAScalar _ _ _ space _ =   error $ "Cannot write to '" ++ space ++ "' memory space."

@@ -205,8 +205,8 @@                                        0, NULL, $exp:(profilingEvent copyScalarFromDev)));               |]   GC.stm-    [C.cstm|if (ctx->failure_is_an_option &&-                     futhark_context_sync(ctx) != 0) { return 1; }|]+    [C.cstm|if (ctx->failure_is_an_option && futhark_context_sync(ctx) != 0)+            { return 1; }|]   pure [C.cexp|$id:val|] readOpenCLScalar _ _ _ space _ =   error $ "Cannot read from '" ++ space ++ "' memory space."

@@ -56,10 +56,10 @@ import Futhark.CodeGen.Backends.GenericPython.Options import qualified Futhark.CodeGen.ImpCode as Imp import Futhark.CodeGen.RTS.Python-import Futhark.Compiler.CLI (CompilerMode (..))+import Futhark.Compiler.Config (CompilerMode (..)) import Futhark.IR.Primitive hiding (Bool) import Futhark.IR.Prop (isBuiltInFunction, subExpVars)-import Futhark.IR.Syntax (Space (..))+import Futhark.IR.Syntax.Core (Space (..)) import Futhark.MonadFreshNames import Futhark.Util (zEncodeString) import Futhark.Util.Pretty (pretty, prettyText)

@@ -365,7 +365,9 @@     declared x = go declared x      set (SetMem x y _) = namesFromList [x, y]-    set (Call _ _ args) = foldMap onArg args+    set (Call dests _ args) =+      -- Some of the dests might not be memory, but it does not matter.+      namesFromList dests <> foldMap onArg args       where         onArg ExpArg {} = mempty         onArg (MemArg x) = oneName x

@@ -143,6 +143,12 @@   mkTV (patElemName pe) int32 <-- sExt32 num_groups opCompiler dest (Inner (SegOp op)) =   segOpCompiler dest op+opCompiler (Pat pes) (Inner (GPUBody _ (Body _ stms res))) = do+  tid <- newVName "tid"+  sKernelThread "gpuseq" tid (defKernelAttrs 1 1) $+    compileStms (freeIn res) stms $+      forM_ (zip pes res) $ \(pe, SubExpRes _ se) ->+        copyDWIM (patElemName pe) [DimFix 0] se [] opCompiler pat e =   compilerBugS $     "opCompiler: Invalid pattern\n  "

@@ -55,8 +55,8 @@     -- to perform well in practice.     SizeHeuristic "" DeviceGPU NumGroups $ 4 * max_compute_units,     SizeHeuristic "" DeviceGPU GroupSize 256,-    SizeHeuristic "" DeviceGPU TileSize 32,-    SizeHeuristic "" DeviceGPU RegTileSize 2,+    SizeHeuristic "" DeviceGPU TileSize 16,+    SizeHeuristic "" DeviceGPU RegTileSize 4,     SizeHeuristic "" DeviceGPU Threshold $ 32 * 1024,     SizeHeuristic "" DeviceCPU LockstepWidth 1,     SizeHeuristic "" DeviceCPU NumGroups max_compute_units,

@@ -3,7 +3,6 @@ -- | Code snippets used by the C backends. module Futhark.CodeGen.RTS.C   ( atomicsH,-    chaselevH,     cudaH,     freeListH,     halfH,@@ -32,11 +31,6 @@ atomicsH :: T.Text atomicsH = $(embedStringFile "rts/c/atomics.h") {-# NOINLINE atomicsH #-}---- | @rts/c/chaselev.h@-chaselevH :: T.Text-chaselevH = $(embedStringFile "rts/c/chaselev.h")-{-# NOINLINE chaselevH #-}  -- | @rts/c/cuda.h@ cudaH :: T.Text

@@ -22,8 +22,9 @@ import Control.Monad import Control.Monad.Except import Data.Bifunctor (first)+import Data.List (sortOn) import qualified Data.List.NonEmpty as NE-import Data.Loc (Loc (NoLoc))+import Data.Loc (Loc (..), posCoff, posFile) import qualified Data.Text.IO as T import qualified Futhark.Analysis.Alias as Alias import Futhark.Compiler.Config@@ -34,7 +35,7 @@ import Futhark.Internalise import Futhark.MonadFreshNames import Futhark.Pipeline-import Futhark.Util.Console (inRed)+import Futhark.Util.Console (inRed, inYellow) import Futhark.Util.Log import Futhark.Util.Pretty (Doc, line, ppr, prettyText, punctuate, stack, text, (</>)) import qualified Language.Futhark as E@@ -135,12 +136,18 @@  -- | Prettyprint program errors as suitable for showing on a text console. pprProgErrors :: NE.NonEmpty ProgError -> Doc-pprProgErrors = stack . punctuate line . map onError . NE.toList+pprProgErrors = stack . punctuate line . map onError . sortOn (rep . locOf) . NE.toList   where+    rep NoLoc = ("", 0)+    rep (Loc p _) = (posFile p, posCoff p)     onError (ProgError NoLoc msg) =       msg     onError (ProgError loc msg) =       text (inRed $ "Error at " <> locStr (srclocOf loc) <> ":") </> msg+    onError (ProgWarning NoLoc msg) =+      msg+    onError (ProgWarning loc msg) =+      text (inYellow $ "Warning at " <> locStr (srclocOf loc) <> ":") </> msg  -- | Throw an exception formatted with 'pprProgErrors' if there's -- an error.

@@ -172,13 +172,6 @@     compilerEntryPoints :: [Name]   } --- | Are we compiling a library or an executable?-data CompilerMode-  = ToLibrary-  | ToExecutable-  | ToServer-  deriving (Eq, Ord, Show)- -- | The configuration of the compiler. newCompilerConfig :: cfg -> CompilerConfig cfg newCompilerConfig x =

@@ -3,11 +3,28 @@   ( FutharkConfig (..),     newFutharkConfig,     Verbosity (..),+    CompilerMode (..),   ) where -import Futhark.IR-import Futhark.Pipeline+import Futhark.IR.Syntax.Core (Name)++-- | Are we compiling a library or an executable?+data CompilerMode+  = ToLibrary+  | ToExecutable+  | ToServer+  deriving (Eq, Ord, Show)++-- | How much information to print to stderr while the compiler is running.+data Verbosity+  = -- | Silence is golden.+    NotVerbose+  | -- | Print messages about which pass is running.+    Verbose+  | -- | Also print logs from individual passes.+    VeryVerbose+  deriving (Eq, Ord)  -- | The compiler configuration.  This only contains options related -- to core compiler functionality, such as reading the initial program

@@ -15,8 +15,10 @@     noLoadedProg,     lpImports,     lpWarnings,+    lpFilePaths,     reloadProg,     extendProg,+    VFS,   ) where @@ -35,7 +37,7 @@ import Data.Bifunctor (first) import Data.List (intercalate, sort) import qualified Data.List.NonEmpty as NE-import Data.Loc (Loc (..), locOf)+import Data.Loc (Loc (..), Located, locOf) import qualified Data.Map as M import Data.Maybe (mapMaybe) import qualified Data.Text as T@@ -43,7 +45,7 @@ import Data.Time.Clock (UTCTime, getCurrentTime) import Futhark.FreshNames import Futhark.Util (interactWithFileSafely, nubOrd, startupTime)-import Futhark.Util.Pretty (Doc, align, line, ppr, text, (</>))+import Futhark.Util.Pretty (Doc, align, ppr, text) import qualified Language.Futhark as E import Language.Futhark.Parser (SyntaxError (..), parseFuthark) import Language.Futhark.Prelude@@ -66,10 +68,18 @@  -- | Note that the location may be 'NoLoc'.  This essentially only -- happens when the problem is that a root file cannot be found.-data ProgError = ProgError Loc Doc+data ProgError+  = ProgError Loc Doc+  | -- | Not actually an error, but we want them reported+    -- with errors.+    ProgWarning Loc Doc  type WithErrors = Either (NE.NonEmpty ProgError) +instance Located ProgError where+  locOf (ProgError l _) = l+  locOf (ProgWarning l _) = l+ -- | A mapping from absolute pathnames to text representing a virtual -- file system.  Before loading a file from the file system, this -- mapping is consulted.  If the desired pathname has an entry here,@@ -279,10 +289,10 @@       case E.checkProg (asImports imports) src import_name prog' of         (prog_ws, Left (E.TypeError loc notes msg)) -> do           let err' = msg <> ppr notes-          Left . singleError . ProgError (locOf loc) $-            if anyWarnings prog_ws-              then ppr prog_ws </> line <> ppr err'-              else ppr err'+              warningToError (wloc, wmsg) = ProgWarning (locOf wloc) wmsg+          Left $+            ProgError (locOf loc) err'+              NE.:| map warningToError (listWarnings prog_ws)         (prog_ws, Right (m, src')) ->           let warnHole (loc, t) =                 singleWarning (E.srclocOf loc) $ "Hole of type: " <> align (ppr t)@@ -345,22 +355,31 @@ lpWarnings :: LoadedProg -> Warnings lpWarnings = foldMap (cfWarnings . lfMod) . lpFiles +-- | The absolute paths of the files that are part of this program.+lpFilePaths :: LoadedProg -> [FilePath]+lpFilePaths = map lfPath . lpFiles+ unchangedImports ::   MonadIO m =>   VNameSource ->+  VFS ->   [LoadedFile CheckedFile] ->   m ([LoadedFile CheckedFile], VNameSource)-unchangedImports src [] = pure ([], src)-unchangedImports src (f : fs)+unchangedImports src _ [] = pure ([], src)+unchangedImports src vfs (f : fs)   | isBuiltin (includeToFilePath (lfImportName f)) =-      first (f :) <$> unchangedImports src fs+      first (f :) <$> unchangedImports src vfs fs   | otherwise = do-      changed <--        maybe True (either (const True) (> lfModTime f))-          <$> liftIO (interactWithFileSafely (getModificationTime $ lfPath f))-      if changed+      let file_path = lfPath f+      if M.member file_path vfs         then pure ([], cfNameSource $ lfMod f)-        else first (f :) <$> unchangedImports src fs+        else do+          changed <-+            maybe True (either (const True) (> lfModTime f))+              <$> liftIO (interactWithFileSafely (getModificationTime file_path))+          if changed+            then pure ([], cfNameSource $ lfMod f)+            else first (f :) <$> unchangedImports src vfs fs  -- | A "loaded program" containing no actual files.  Use this as a -- starting point for 'reloadProg'@@ -375,10 +394,10 @@ -- | Find out how many of the old imports can be used.  Here we are -- forced to be overly conservative, because our type checker -- enforces a linear ordering.-usableLoadedProg :: MonadIO m => LoadedProg -> [FilePath] -> m LoadedProg-usableLoadedProg (LoadedProg roots imports src) new_roots+usableLoadedProg :: MonadIO m => LoadedProg -> VFS -> [FilePath] -> m LoadedProg+usableLoadedProg (LoadedProg roots imports src) vfs new_roots   | sort roots == sort new_roots = do-      (imports', src') <- unchangedImports src imports+      (imports', src') <- unchangedImports src vfs imports       pure $ LoadedProg [] imports' src'   | otherwise =       pure noLoadedProg@@ -407,7 +426,7 @@   VFS ->   IO (Either (NE.NonEmpty ProgError) LoadedProg) reloadProg lp new_roots vfs = do-  lp' <- usableLoadedProg lp new_roots+  lp' <- usableLoadedProg lp vfs new_roots   extendProg lp' new_roots vfs  -- | Read and type-check some Futhark files.

@@ -85,6 +85,7 @@     eSliceArray,     eBlank,     eAll,+    eDimInBounds,     eOutOfBounds,      -- * Other building blocks@@ -374,6 +375,14 @@   pure $ BasicOp $ Index arr $ fullSlice arr_t $ skips ++ [slice i' n']   where     slice j m = DimSlice j m (constant (1 :: Int64))++-- | @eInBoundsForDim w i@ produces @0 <= i < w@.+eDimInBounds :: MonadBuilder m => m (Exp (Rep m)) -> m (Exp (Rep m)) -> m (Exp (Rep m))+eDimInBounds w i =+  eBinOp+    LogAnd+    (eCmpOp (CmpSle Int64) (eSubExp (intConst Int64 0)) i)+    (eCmpOp (CmpSlt Int64) i w)  -- | Are these indexes out-of-bounds for the array? eOutOfBounds ::

@@ -141,12 +141,13 @@     forM imports $ \(current, fm) ->       let ctx =             Context-              current-              fm-              imports-              mempty-              file_map-              (progModuleTypes $ fileProg fm)+              { ctxCurrent = current,+                ctxFileMod = fm,+                ctxImports = imports,+                ctxNoLink = mempty,+                ctxFileMap = file_map,+                ctxVisibleMTys = progModuleTypes $ fileProg fm+              }        in flip runReaderT ctx $ do             (first_paragraph, maybe_abstract, maybe_sections) <- headerDoc $ fileProg fm @@ -519,10 +520,10 @@  prettyShapeDecl :: ShapeDecl (DimDecl VName) -> DocM Html prettyShapeDecl (ShapeDecl ds) =-  mconcat <$> mapM (fmap brackets . dimDeclHtml) ds+  mconcat <$> mapM dimDeclHtml ds  typeArgHtml :: TypeArg (DimDecl VName) -> DocM Html-typeArgHtml (TypeArgDim d _) = brackets <$> dimDeclHtml d+typeArgHtml (TypeArgDim d _) = dimDeclHtml d typeArgHtml (TypeArgType t _) = typeHtml t  modParamHtml :: [ModParamBase Info VName] -> DocM Html@@ -616,7 +617,7 @@   TEArray at d _ -> do     at' <- typeExpHtml at     d' <- dimExpHtml d-    pure $ brackets d' <> at'+    pure $ d' <> at'   TETuple ts _ -> parens . commas <$> mapM typeExpHtml ts   TERecord fs _ -> braces . commas <$> mapM ppField fs     where@@ -691,14 +692,14 @@   concat (replicate (length (splitPath src) - 1) "../") ++ dest  dimDeclHtml :: DimDecl VName -> DocM Html-dimDeclHtml (NamedDim v) = qualNameHtml v-dimDeclHtml (ConstDim n) = pure $ toHtml (show n)-dimDeclHtml AnyDim {} = pure mempty+dimDeclHtml (NamedDim v) = brackets <$> qualNameHtml v+dimDeclHtml (ConstDim n) = pure $ brackets $ toHtml (show n)+dimDeclHtml AnyDim {} = pure $ brackets mempty  dimExpHtml :: DimExp VName -> DocM Html-dimExpHtml DimExpAny = pure mempty-dimExpHtml (DimExpNamed v _) = qualNameHtml v-dimExpHtml (DimExpConst n _) = pure $ toHtml (show n)+dimExpHtml DimExpAny = pure $ brackets mempty+dimExpHtml (DimExpNamed v _) = brackets <$> qualNameHtml v+dimExpHtml (DimExpConst n _) = pure $ brackets $ toHtml (show n)  typeArgExpHtml :: TypeArgExp VName -> DocM Html typeArgExpHtml (TypeArgExpDim d _) = dimExpHtml d

@@ -38,6 +38,7 @@     removeFunDefAliases,     removeExpAliases,     removeStmAliases,+    removeBodyAliases,     removeLambdaAliases,     removePatAliases,     removeScopeAliases,@@ -46,6 +47,7 @@     AliasesAndConsumed,     trackAliases,     mkStmsAliases,+    consumedInStms,   ) where @@ -226,6 +228,13 @@   Stm rep removeStmAliases = runIdentity . rephraseStm removeAliases +-- | Remove alias information from body.+removeBodyAliases ::+  CanBeAliased (Op rep) =>+  Body (Aliases rep) ->+  Body rep+removeBodyAliases = runIdentity . rephraseBody removeAliases+ -- | Remove alias information from lambda. removeLambdaAliases ::   CanBeAliased (Op rep) =>@@ -320,6 +329,10 @@     Names   ) +-- | The variables consumed in these statements.+consumedInStms :: Aliased rep => Stms rep -> Names+consumedInStms = snd . flip mkStmsAliases []+ -- | A helper function for computing the aliases of a sequence of -- statements.  You'd use this while recursing down the statements -- from first to last.  The 'AliasesAndConsumed' parameter is the@@ -348,13 +361,13 @@     aliasesOfAliases = mconcat . map look . namesToList     look k = M.findWithDefault mempty k aliasmap -mkAliasedLetStm ::+mkAliasedStm ::   (ASTRep rep, CanBeAliased (Op rep)) =>   Pat (LetDec rep) ->   StmAux (ExpDec rep) ->   Exp (Aliases rep) ->   Stm (Aliases rep)-mkAliasedLetStm pat (StmAux cs attrs dec) e =+mkAliasedStm pat (StmAux cs attrs dec) e =   Let     (mkAliasedPat pat e)     (StmAux cs attrs (AliasDec $ consumedInExp e, dec))@@ -372,7 +385,7 @@     env <- asksScope removeScopeAliases     flip runReaderT env $ do       Let pat dec _ <- mkLetNames names $ removeExpAliases e-      pure $ mkAliasedLetStm pat dec e+      pure $ mkAliasedStm pat dec e    mkBody stms res =     let Body bodyrep _ _ = mkBody (fmap removeStmAliases stms) res

@@ -25,10 +25,13 @@   ) where +import Control.Monad+import qualified Data.Sequence as SQ+import qualified Futhark.Analysis.Alias as Alias import Futhark.Analysis.Metrics import qualified Futhark.Analysis.SymbolTable as ST import Futhark.IR-import Futhark.IR.Aliases (Aliases)+import Futhark.IR.Aliases (Aliases, removeBodyAliases) import Futhark.IR.GPU.Sizes import Futhark.IR.Prop.Aliases import Futhark.IR.SegOp@@ -45,7 +48,6 @@     (<+>),   ) import qualified Futhark.Util.Pretty as PP-import Prelude hiding (id, (.))  -- | At which level the *body* of a t'SegOp' executes. data SegLevel@@ -240,6 +242,9 @@     SegOp (SegOp SegLevel rep)   | SizeOp SizeOp   | OtherOp op+  | -- | Code to run sequentially on the GPU,+    -- in a single thread.+    GPUBody [Type] (Body rep)   deriving (Eq, Ord, Show)  -- | A helper for defining 'TraverseOpStms'.@@ -250,6 +255,9 @@ traverseHostOpStms _ f (SegOp segop) = SegOp <$> traverseSegOpStms f segop traverseHostOpStms _ _ (SizeOp sizeop) = pure $ SizeOp sizeop traverseHostOpStms onOtherOp f (OtherOp other) = OtherOp <$> onOtherOp f other+traverseHostOpStms _ f (GPUBody ts body) = do+  stms <- f mempty $ bodyStms body+  pure $ GPUBody ts $ body {bodyStms = stms}  instance (ASTRep rep, Substitute op) => Substitute (HostOp rep op) where   substituteNames substs (SegOp op) =@@ -258,50 +266,65 @@     OtherOp $ substituteNames substs op   substituteNames substs (SizeOp op) =     SizeOp $ substituteNames substs op+  substituteNames substs (GPUBody ts body) =+    GPUBody (substituteNames substs ts) (substituteNames substs body)  instance (ASTRep rep, Rename op) => Rename (HostOp rep op) where   rename (SegOp op) = SegOp <$> rename op   rename (OtherOp op) = OtherOp <$> rename op   rename (SizeOp op) = SizeOp <$> rename op+  rename (GPUBody ts body) = GPUBody <$> rename ts <*> rename body  instance (ASTRep rep, IsOp op) => IsOp (HostOp rep op) where   safeOp (SegOp op) = safeOp op   safeOp (OtherOp op) = safeOp op   safeOp (SizeOp op) = safeOp op+  safeOp GPUBody {} = True    cheapOp (SegOp op) = cheapOp op   cheapOp (OtherOp op) = cheapOp op   cheapOp (SizeOp op) = cheapOp op+  cheapOp (GPUBody types body) =+    -- Current GPUBody usage only benefits from hoisting kernels that+    -- transfer scalars to device.+    SQ.null (bodyStms body) && all ((== 0) . arrayRank) types  instance TypedOp op => TypedOp (HostOp rep op) where   opType (SegOp op) = opType op   opType (OtherOp op) = opType op   opType (SizeOp op) = opType op+  opType (GPUBody ts _) =+    pure $ staticShapes $ map (`arrayOfRow` intConst Int64 1) ts  instance (Aliased rep, AliasedOp op, ASTRep rep) => AliasedOp (HostOp rep op) where   opAliases (SegOp op) = opAliases op   opAliases (OtherOp op) = opAliases op   opAliases (SizeOp op) = opAliases op+  opAliases (GPUBody ts _) = map (const mempty) ts    consumedInOp (SegOp op) = consumedInOp op   consumedInOp (OtherOp op) = consumedInOp op   consumedInOp (SizeOp op) = consumedInOp op+  consumedInOp (GPUBody _ body) = consumedInBody body  instance (ASTRep rep, FreeIn op) => FreeIn (HostOp rep op) where   freeIn' (SegOp op) = freeIn' op   freeIn' (OtherOp op) = freeIn' op   freeIn' (SizeOp op) = freeIn' op+  freeIn' (GPUBody ts body) = freeIn' ts <> freeIn' body  instance (CanBeAliased (Op rep), CanBeAliased op, ASTRep rep) => CanBeAliased (HostOp rep op) where   type OpWithAliases (HostOp rep op) = HostOp (Aliases rep) (OpWithAliases op)    addOpAliases aliases (SegOp op) = SegOp $ addOpAliases aliases op+  addOpAliases aliases (GPUBody ts body) = GPUBody ts $ Alias.analyseBody aliases body   addOpAliases aliases (OtherOp op) = OtherOp $ addOpAliases aliases op   addOpAliases _ (SizeOp op) = SizeOp op    removeOpAliases (SegOp op) = SegOp $ removeOpAliases op   removeOpAliases (OtherOp op) = OtherOp $ removeOpAliases op   removeOpAliases (SizeOp op) = SizeOp op+  removeOpAliases (GPUBody ts body) = GPUBody ts $ removeBodyAliases body  instance (CanBeWise (Op rep), CanBeWise op, ASTRep rep) => CanBeWise (HostOp rep op) where   type OpWithWisdom (HostOp rep op) = HostOp (Wise rep) (OpWithWisdom op)@@ -309,10 +332,12 @@   removeOpWisdom (SegOp op) = SegOp $ removeOpWisdom op   removeOpWisdom (OtherOp op) = OtherOp $ removeOpWisdom op   removeOpWisdom (SizeOp op) = SizeOp op+  removeOpWisdom (GPUBody ts body) = GPUBody ts $ removeBodyWisdom body    addOpWisdom (SegOp op) = SegOp $ addOpWisdom op   addOpWisdom (OtherOp op) = OtherOp $ addOpWisdom op   addOpWisdom (SizeOp op) = SizeOp op+  addOpWisdom (GPUBody ts body) = GPUBody ts $ informBody body  instance (ASTRep rep, ST.IndexOp op) => ST.IndexOp (HostOp rep op) where   indexOp vtable k (SegOp op) is = ST.indexOp vtable k op is@@ -323,11 +348,14 @@   ppr (SegOp op) = ppr op   ppr (OtherOp op) = ppr op   ppr (SizeOp op) = ppr op+  ppr (GPUBody ts body) =+    "gpu" <+> PP.colon <+> ppTuple' ts <+> PP.nestedBlock "{" "}" (ppr body)  instance (OpMetrics (Op rep), OpMetrics op) => OpMetrics (HostOp rep op) where   opMetrics (SegOp op) = opMetrics op   opMetrics (OtherOp op) = opMetrics op   opMetrics (SizeOp op) = opMetrics op+  opMetrics (GPUBody _ body) = inside "GPUBody" $ bodyMetrics body  checkSegLevel ::   TC.Checkable rep =>@@ -356,5 +384,18 @@ typeCheckHostOp checker lvl _ (SegOp op) =   TC.checkOpWith (checker $ segLevel op) $     typeCheckSegOp (checkSegLevel lvl) op+typeCheckHostOp _ Just {} _ GPUBody {} =+  TC.bad $ TC.TypeError "GPUBody may not be nested in SegOps." typeCheckHostOp _ _ f (OtherOp op) = f op typeCheckHostOp _ _ _ (SizeOp op) = typeCheckSizeOp op+typeCheckHostOp _ Nothing _ (GPUBody ts body) = do+  mapM_ TC.checkType ts+  void $ TC.checkBody body+  body_ts <-+    extendedScope+      (traverse subExpResType (bodyResult body))+      (scopeOf (bodyStms body))+  unless (body_ts == ts) . TC.bad . TC.TypeError . unlines $+    [ "Expected type: " ++ prettyTuple ts,+      "Got body type: " ++ prettyTuple body_ts+    ]

@@ -15,6 +15,7 @@   ) where +import qualified Futhark.Analysis.UsageTable as UT import Futhark.IR.GPU import qualified Futhark.IR.SOACS.Simplify as SOAC import Futhark.MonadFreshNames@@ -72,6 +73,17 @@ simplifyKernelOp _ (SizeOp (CalcNumGroups w max_num_groups group_size)) = do   w' <- Engine.simplify w   pure (SizeOp $ CalcNumGroups w' max_num_groups group_size, mempty)+simplifyKernelOp _ (GPUBody ts body) = do+  ts' <- Engine.simplify ts+  (hoisted, body') <-+    Engine.simplifyBody keepOnGPU mempty (map (const mempty) ts) body+  pure (GPUBody ts' body', hoisted)+  where+    keepOnGPU _ _ = keepExpOnGPU . stmExp+    keepExpOnGPU (BasicOp Index {}) = True+    keepExpOnGPU (BasicOp (ArrayLit _ t)) | primType t = True+    keepExpOnGPU DoLoop {} = True+    keepExpOnGPU _ = False  instance TraverseOpStms (Wise GPU) where   traverseOpStms = traverseHostOpStms traverseSOACStms@@ -96,7 +108,28 @@       [ RuleOp SOAC.simplifyKnownIterationSOAC,         RuleOp SOAC.removeReplicateMapping,         RuleOp SOAC.liftIdentityMapping,-        RuleOp SOAC.simplifyMapIota+        RuleOp SOAC.simplifyMapIota,+        RuleOp SOAC.removeUnusedSOACInput       ]-      [ RuleBasicOp removeUnnecessaryCopy+      [ RuleBasicOp removeUnnecessaryCopy,+        RuleOp removeDeadGPUBodyResult       ]++-- | Remove the unused return values of a GPUBody.+removeDeadGPUBodyResult :: BottomUpRuleOp (Wise GPU)+removeDeadGPUBodyResult (_, used) pat aux (GPUBody types body)+  | -- Figure out which of the names in 'pat' are used...+    pat_used <- map (`UT.isUsedDirectly` used) $ patNames pat,+    -- If they are not all used, then this rule applies.+    not (and pat_used) =+      -- Remove the parts of the GPUBody results that correspond to dead+      -- return value bindings.  Note that this leaves dead code in the+      -- kernel, but that will be removed later.+      let pick :: [a] -> [a]+          pick = map snd . filter fst . zip pat_used+          pat' = pick (patElems pat)+          types' = pick types+          body' = body {bodyResult = pick (bodyResult body)}+       in Simplify $ auxing aux $ letBind (Pat pat') $ Op $ GPUBody types' body'+  | otherwise = Skip+removeDeadGPUBodyResult _ _ _ _ = Skip

@@ -143,7 +143,7 @@                 v_copy <- newVName $ baseString v <> "_nonext_copy"                 let v_pat =                       Pat [PatElem v_copy $ MemArray pt shape u $ ArrayIn mem ixfun]-                addStm $ mkWiseLetStm v_pat (defAux ()) $ BasicOp (Copy v)+                addStm $ mkWiseStm v_pat (defAux ()) $ BasicOp (Copy v)                 pure $ SubExpRes cs $ Var v_copy             | Just mem <- lookup (patElemName pat_elem) oldmem_to_mem =                 pure $ SubExpRes cs $ Var mem

@@ -578,6 +578,8 @@       keyword "redomap" *> pScrema pRedomapForm,       keyword "scanomap" *> pScrema pScanomapForm,       keyword "screma" *> pScrema pScremaForm,+      keyword "vjp" *> pVJP,+      keyword "jvp" *> pJVP,       pScatter,       pHist,       pStream@@ -659,6 +661,18 @@           <*> pure SOAC.Sequential           <*> braces (pSubExp `sepBy` pComma) <* pComma           <*> pLambda pr+    pVJP =+      parens $+        SOAC.VJP+          <$> pLambda pr <* pComma+          <*> braces (pSubExp `sepBy` pComma) <* pComma+          <*> braces (pSubExp `sepBy` pComma)+    pJVP =+      parens $+        SOAC.JVP+          <$> pLambda pr <* pComma+          <*> braces (pSubExp `sepBy` pComma) <* pComma+          <*> braces (pSubExp `sepBy` pComma)  pSizeClass :: Parser GPU.SizeClass pSizeClass =@@ -843,7 +857,8 @@   choice     [ GPU.SegOp <$> pSegOp pr pSegLevel,       GPU.SizeOp <$> pSizeOp,-      GPU.OtherOp <$> pOther+      GPU.OtherOp <$> pOther,+      keyword "gpu" $> GPU.GPUBody <*> (pColon *> pTypes) <*> braces (pBody pr)     ]  pMCOp :: PR rep -> Parser op -> Parser (MC.MCOp rep op)

@@ -1670,7 +1670,7 @@ negativeIshInt (Int32Value k) = k < 0 negativeIshInt (Int64Value k) = k < 0 --- | The size of a value of a given primitive type in bites.+-- | The size of a value of a given primitive type in bits. primBitSize :: PrimType -> Int primBitSize = (* 8) . primByteSize 

@@ -52,7 +52,7 @@ vnameAliases :: VName -> Names vnameAliases = oneName --- | The alises of a subexpression.+-- | The aliases of a subexpression. subExpAliases :: SubExp -> Names subExpAliases Constant {} = mempty subExpAliases (Var v) = vnameAliases v

@@ -10,6 +10,7 @@   ( -- * Free names     Names,     namesIntMap,+    namesIntSet,     nameIn,     notNameIn,     oneName,@@ -46,6 +47,7 @@ import Control.Monad.State.Strict import Data.Foldable import qualified Data.IntMap.Strict as IM+import qualified Data.IntSet as IS import qualified Data.Map.Strict as M import Futhark.IR.Prop.Patterns import Futhark.IR.Prop.Scope@@ -62,6 +64,10 @@ -- | Retrieve the data structure underlying the names representation. namesIntMap :: Names -> IM.IntMap VName namesIntMap (Names m) = m++-- | Retrieve the set of tags in the names set.+namesIntSet :: Names -> IS.IntSet+namesIntSet (Names m) = IM.keysSet m  instance Ord Names where   x `compare` y = if x == y then EQ else LT

@@ -13,6 +13,7 @@     staticShapes,     staticShapes1,     primType,+    isAcc,     arrayOf,     arrayOfRow,     arrayOfShape,@@ -137,7 +138,7 @@ unique = (== Unique) . uniqueness  -- | Convert types with non-existential shapes to types with--- non-existential shapes.  Only the representation is changed, so all+-- existential shapes.  Only the representation is changed, so all -- the shapes will be 'Free'. staticShapes :: [TypeBase Shape u] -> [TypeBase ExtShape u] staticShapes = map staticShapes1@@ -287,12 +288,17 @@ primType Prim {} = True primType _ = False +-- | Is this an accumulator?+isAcc :: TypeBase shape u -> Bool+isAcc Acc {} = True+isAcc _ = False+ -- | Returns the bottommost type of an array.  For @[][]i32@, this -- would be @i32@.  If the given type is not an array, it is returned. elemType :: TypeBase shape u -> PrimType elemType (Array t _ _) = t elemType (Prim t) = t-elemType Acc {} = error "Acc"+elemType Acc {} = error "elemType Acc" elemType Mem {} = error "elemType Mem"  -- | Swap the two outer dimensions of the type.

@@ -4,6 +4,7 @@ -- | A simple representation with SOACs and nested parallelism. module Futhark.IR.SOACS   ( SOACS,+    usesAD,      -- * Module re-exports     module Futhark.IR.Prop,@@ -46,3 +47,27 @@ instance BuilderOps SOACS  instance PrettyRep SOACS++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 = bodyUsesAD . lambdaBody+    expUsesAD (Op JVP {}) = True+    expUsesAD (Op VJP {}) = True+    expUsesAD (Op (Stream _ _ _ _ lam)) = lamUsesAD lam+    expUsesAD (Op (Screma _ _ (ScremaForm scans reds lam))) =+      lamUsesAD lam+        || any (lamUsesAD . scanLambda) scans+        || any (lamUsesAD . redLambda) reds+    expUsesAD (Op (Hist _ _ ops lam)) =+      lamUsesAD lam || any (lamUsesAD . histOp) ops+    expUsesAD (Op (Scatter _ _ lam _)) =+      lamUsesAD lam+    expUsesAD (If _ tbody fbody _) = bodyUsesAD tbody || bodyUsesAD fbody+    expUsesAD (DoLoop _ _ body) = bodyUsesAD body+    expUsesAD (WithAcc _ lam) = lamUsesAD lam+    expUsesAD BasicOp {} = False+    expUsesAD Apply {} = False

@@ -122,6 +122,10 @@     --     -- The final lambda produces indexes and values for the 'HistOp's.     Hist SubExp [VName] [HistOp rep] (Lambda rep)+  | -- FIXME: this should not be here+    JVP (Lambda rep) [SubExp] [SubExp]+  | -- FIXME: this should not be here+    VJP (Lambda rep) [SubExp] [SubExp]   | -- | A combination of scan, reduction, and map.  The first     -- t'SubExp' is the size of the input arrays.     Screma SubExp [VName] (ScremaForm rep)@@ -400,6 +404,14 @@   SOACMapper frep trep m ->   SOAC frep ->   m (SOAC trep)+mapSOACM tv (JVP lam args vec) =+  JVP <$> mapOnSOACLambda tv lam+    <*> mapM (mapOnSOACSubExp tv) args+    <*> mapM (mapOnSOACSubExp tv) vec+mapSOACM tv (VJP lam args vec) =+  VJP <$> mapOnSOACLambda tv lam+    <*> mapM (mapOnSOACSubExp tv) args+    <*> mapM (mapOnSOACSubExp tv) vec mapSOACM tv (Stream size arrs form accs lam) =   Stream <$> mapOnSOACSubExp tv size     <*> mapM (mapOnSOACVName tv) arrs@@ -490,7 +502,13 @@       renamer = SOACMapper rename rename rename  -- | The type of a SOAC.-soacType :: SOAC rep -> [Type]+soacType :: Typed (LParamInfo rep) => SOAC rep -> [Type]+soacType (JVP lam _ _) =+  lambdaReturnType lam+    ++ lambdaReturnType lam+soacType (VJP lam _ _) =+  lambdaReturnType lam+    ++ map paramType (lambdaParams lam) soacType (Stream outersize _ _ accs lam) =   map (substNamesInType substs) rtp   where@@ -509,12 +527,14 @@ soacType (Screma w _arrs form) =   scremaType w form -instance TypedOp (SOAC rep) where+instance ASTRep rep => TypedOp (SOAC rep) where   opType = pure . staticShapes . soacType  instance (ASTRep rep, Aliased rep) => AliasedOp (SOAC rep) where   opAliases = map (const mempty) . soacType +  consumedInOp JVP {} = mempty+  consumedInOp VJP {} = mempty   -- Only map functions can consume anything.  The operands to scan   -- and reduce functions are always considered "fresh".   consumedInOp (Screma _ arrs (ScremaForm _ _ map_lam)) =@@ -556,6 +576,10 @@   where   type OpWithAliases (SOAC rep) = SOAC (Aliases rep) +  addOpAliases aliases (JVP lam args vec) =+    JVP (Alias.analyseLambda aliases lam) args vec+  addOpAliases aliases (VJP lam args vec) =+    VJP (Alias.analyseLambda aliases lam) args vec   addOpAliases aliases (Stream size arr form accs lam) =     Stream size arr (analyseStreamForm form) accs $       Alias.analyseLambda aliases lam@@ -648,6 +672,8 @@  -- | Type-check a SOAC. typeCheckSOAC :: TC.Checkable rep => SOAC (Aliases rep) -> TC.TypeM rep ()+typeCheckSOAC JVP {} = pure ()+typeCheckSOAC VJP {} = pure () typeCheckSOAC (Stream size arrexps form accexps lam) = do   TC.require [Prim int64] size   accargs <- mapM TC.checkArg accexps@@ -820,6 +846,10 @@           ++ " wrong for given scan and reduction functions."  instance OpMetrics (Op rep) => OpMetrics (SOAC rep) where+  opMetrics (VJP lam _ _) =+    inside "VJP" $ lambdaMetrics lam+  opMetrics (JVP lam _ _) =+    inside "JVP" $ lambdaMetrics lam   opMetrics (Stream _ _ _ _ lam) =     inside "Stream" $ lambdaMetrics lam   opMetrics (Scatter _len _ lam _) =@@ -833,6 +863,22 @@       lambdaMetrics map_lam  instance PrettyRep rep => PP.Pretty (SOAC rep) where+  ppr (VJP lam args vec) =+    text "vjp"+      <> parens+        ( PP.align $+            ppr lam <> comma+              </> PP.braces (commasep $ map ppr args) <> comma+              </> PP.braces (commasep $ map ppr vec)+        )+  ppr (JVP lam args vec) =+    text "jvp"+      <> parens+        ( PP.align $+            ppr lam <> comma+              </> PP.braces (commasep $ map ppr args) <> comma+              </> PP.braces (commasep $ map ppr vec)+        )   ppr (Stream size arrs form acc lam) =     case form of       Parallel o comm lam0 ->

@@ -17,6 +17,7 @@     HasSOAC (..),     simplifyKnownIterationSOAC,     removeReplicateMapping,+    removeUnusedSOACInput,     liftIdentityMapping,     simplifyMapIota,     SOACS,@@ -85,6 +86,16 @@ simplifySOAC ::   Simplify.SimplifiableRep rep =>   Simplify.SimplifyOp rep (SOAC (Wise rep))+simplifySOAC (VJP lam arr vec) = do+  (lam', hoisted) <- Engine.simplifyLambda lam+  arr' <- mapM Engine.simplify arr+  vec' <- mapM Engine.simplify vec+  pure (VJP lam' arr' vec', hoisted)+simplifySOAC (JVP lam arr vec) = do+  (lam', hoisted) <- Engine.simplifyLambda lam+  arr' <- mapM Engine.simplify arr+  vec' <- mapM Engine.simplify vec+  pure (JVP lam' arr' vec', hoisted) simplifySOAC (Stream outerdim arr form nes lam) = do   outerdim' <- Engine.simplify outerdim   (form', form_hoisted) <- simplifyStreamForm form@@ -392,17 +403,21 @@           Left (p, v)  -- | Remove inputs that are not used inside the SOAC.-removeUnusedSOACInput :: TopDownRuleOp (Wise SOACS)-removeUnusedSOACInput _ pat aux (Screma w arrs (ScremaForm scan reduce map_lam))-  | (used, unused) <- partition usedInput params_and_arrs,+removeUnusedSOACInput ::+  forall rep.+  (Aliased rep, Buildable rep, BuilderOps rep, HasSOAC rep) =>+  TopDownRuleOp rep+removeUnusedSOACInput _ pat aux op+  | Just (Screma w arrs form :: SOAC rep) <- asSOAC op,+    ScremaForm scan reduce map_lam <- form,+    (used, unused) <- partition (usedInput map_lam) (zip (lambdaParams map_lam) arrs),     not (null unused) = Simplify $ do       let (used_params, used_arrs) = unzip used           map_lam' = map_lam {lambdaParams = used_params}-      auxing aux $ letBind pat $ Op $ Screma w used_arrs (ScremaForm scan reduce map_lam')+      auxing aux $ letBind pat $ Op $ soacOp $ Screma w used_arrs (ScremaForm scan reduce map_lam')   where-    params_and_arrs = zip (lambdaParams map_lam) arrs-    used_in_body = freeIn $ lambdaBody map_lam-    usedInput (param, _) = paramName param `nameIn` used_in_body+    used_in_body map_lam = freeIn $ lambdaBody map_lam+    usedInput map_lam (param, _) = paramName param `nameIn` used_in_body map_lam removeUnusedSOACInput _ _ _ _ = Skip  removeDeadMapping :: BottomUpRuleOp (Wise SOACS)

@@ -56,6 +56,7 @@  import Control.Category import Control.Monad.Identity hiding (mapM_)+import Control.Monad.Reader hiding (mapM_) import Control.Monad.State.Strict import Control.Monad.Writer hiding (mapM_) import Data.Bifunctor (first)@@ -1178,6 +1179,8 @@           `Engine.orIf` Engine.isOp           `Engine.orIf` par_blocker           `Engine.orIf` Engine.isConsumed+          `Engine.orIf` Engine.isConsuming+          `Engine.orIf` Engine.isDeviceMigrated    -- Ensure we do not try to use anything that is consumed in the result.   (body_res, body_stms, hoisted) <-@@ -1201,6 +1204,12 @@     consumedInResult _ =       [] +simplifyLambda ::+  Engine.SimplifiableRep rep =>+  Lambda (Wise rep) ->+  Engine.SimpleM rep (Lambda (Wise rep), Stms (Wise rep))+simplifyLambda = Engine.blockMigrated . Engine.simplifyLambda+ segSpaceSymbolTable :: ASTRep rep => SegSpace -> ST.SymbolTable rep segSpaceSymbolTable (SegSpace flat gtids_and_dims) =   foldl' f (ST.fromScope $ M.singleton flat $ IndexName Int64) gtids_and_dims@@ -1214,7 +1223,7 @@ simplifySegBinOp (SegBinOp comm lam nes shape) = do   (lam', hoisted) <-     Engine.localVtable (\vtable -> vtable {ST.simplifyMemory = True}) $-      Engine.simplifyLambda lam+      simplifyLambda lam   shape' <- Engine.simplify shape   nes' <- mapM Engine.simplify nes   pure (SegBinOp comm lam' nes' shape', hoisted)@@ -1276,7 +1285,7 @@         (lam', op_hoisted) <-           Engine.localVtable (<> scope_vtable) $             Engine.localVtable (\vtable -> vtable {ST.simplifyMemory = True}) $-              Engine.simplifyLambda lam+              simplifyLambda lam         pure           ( HistOp w' rf' arrs' nes' dims' lam',             op_hoisted@@ -1301,7 +1310,7 @@  -- | Simplification rules for simplifying 'SegOp's. segOpRules ::-  (HasSegOp rep, BuilderOps rep, Buildable rep) =>+  (HasSegOp rep, BuilderOps rep, Buildable rep, Aliased rep) =>   RuleBook rep segOpRules =   ruleBook [RuleOp segOpRuleTopDown] [RuleOp segOpRuleBottomUp]@@ -1316,7 +1325,7 @@       Skip  segOpRuleBottomUp ::-  (HasSegOp rep, BuilderOps rep) =>+  (HasSegOp rep, BuilderOps rep, Aliased rep) =>   BottomUpRuleOp rep segOpRuleBottomUp vtable pat dec op   | Just op' <- asSegOp op =@@ -1431,7 +1440,7 @@   (kts, body, sum $ map (length . histDest) ops, SegHist lvl space ops)  bottomUpSegOp ::-  (HasSegOp rep, BuilderOps rep) =>+  (Aliased rep, HasSegOp rep, BuilderOps rep) =>   (ST.SymbolTable rep, UT.UsageTable) ->   Pat (LetDec rep) ->   StmAux (ExpDec rep) ->@@ -1454,15 +1463,16 @@     (kpes' == kpes)     cannotSimplify -  kbody <-+  kbody' <-     localScope (scopeOfSegSpace space) $       mkKernelBodyM kstms' kres' -  addStm $ Let (Pat kpes') dec $ Op $ segOp $ mk_segop kts' kbody+  addStm $ Let (Pat kpes') dec $ Op $ segOp $ mk_segop kts' kbody'   where-    (kts, KernelBody _ kstms kres, num_nonmap_results, mk_segop) =+    (kts, kbody@(KernelBody _ kstms kres), num_nonmap_results, mk_segop) =       segOpGuts segop     free_in_kstms = foldMap freeIn kstms+    consumed_in_segop = consumedInKernelBody kbody     space = segSpace segop      sliceWithGtidsFixed stm@@ -1494,6 +1504,7 @@                 letBindNames [patElemName kpe'] . BasicOp . Index arr $                   Slice $ outer_slice <> remaining_slice           if patElemName kpe `UT.isConsumed` used+            || arr `nameIn` consumed_in_segop             then do               precopy <- newVName $ baseString (patElemName kpe) <> "_precopy"               index kpe {patElemName = precopy}

@@ -157,10 +157,12 @@     stmsFromList,     stmsToList,     stmsHead,+    stmsLast,     subExpRes,     subExpsRes,     varRes,     varsRes,+    subExpResVName,   ) where @@ -245,6 +247,12 @@   stm Seq.:< stms' -> Just (stm, stms')   Seq.EmptyL -> Nothing +-- | The last statement in the sequence, if any.+stmsLast :: Stms lore -> Maybe (Stms lore, Stm lore)+stmsLast stms = case Seq.viewr stms of+  stms' Seq.:> stm -> Just (stms', stm)+  Seq.EmptyR -> Nothing+ -- | A pairing of a subexpression and some certificates. data SubExpRes = SubExpRes   { resCerts :: Certs,@@ -268,6 +276,11 @@ varsRes :: [VName] -> Result varsRes = map varRes +-- | The 'VName' of a 'SubExpRes', if it exists.+subExpResVName :: SubExpRes -> Maybe VName+subExpResVName (SubExpRes _ (Var v)) = Just v+subExpResVName _ = Nothing+ -- | The result of a body is a sequence of subexpressions. type Result = [SubExpRes] @@ -363,18 +376,14 @@     Update Safety VName (Slice SubExp) SubExp   | FlatIndex VName (FlatSlice SubExp)   | FlatUpdate VName (FlatSlice SubExp) VName-  | -- | @-    -- concat(0, [1] :| [[2, 3, 4], [5, 6]], 6) = [1, 2, 3, 4, 5, 6]@.-    -- @+  | -- | @concat(0, [1] :| [[2, 3, 4], [5, 6]], 6) = [1, 2, 3, 4, 5, 6]@     --     -- Concatenates the non-empty list of 'VName' resulting in an     -- array of length t'SubExp'. The 'Int' argument is used to     -- specify the dimension along which the arrays are     -- concatenated. For instance:     ---    -- @-    -- concat(1, [[1,2], [3, 4]] :| [[[5,6]], [[7, 8]]], 4) = [[1, 2, 5, 6], [3, 4, 7, 8]]-    -- @+    -- @concat(1, [[1,2], [3, 4]] :| [[[5,6]], [[7, 8]]], 4) = [[1, 2, 5, 6], [3, 4, 7, 8]]@     Concat Int (NonEmpty VName) SubExp   | -- | Copy the given array.  The result will not alias anything.     Copy VName

@@ -43,6 +43,7 @@     oneAttr,     inAttrs,     withoutAttrs,+    mapAttrs,      -- * Values     PrimValue (..),@@ -553,3 +554,7 @@ -- | @x `withoutAttrs` y@ gives @x@ except for any attributes also in @y@. withoutAttrs :: Attrs -> Attrs -> Attrs withoutAttrs (Attrs x) (Attrs y) = Attrs $ x `S.difference` y++-- | Map a function over an attribute set.+mapAttrs :: (Attr -> a) -> Attrs -> [a]+mapAttrs f (Attrs attrs) = map f $ S.toList attrs

@@ -129,9 +129,17 @@ lookupMod' :: QualName VName -> Scope -> Either String Mod lookupMod' mname scope =   let (mname', scope') = lookupSubstInScope mname scope-   in maybe (Left $ bad mname') Right $ M.lookup (qualLeaf mname') $ scopeMods scope'+   in maybe (Left $ bad mname') (Right . extend) $ M.lookup (qualLeaf mname') $ scopeMods scope'   where     bad mname' = "Unknown module: " ++ pretty mname ++ " (" ++ pretty mname' ++ ")"+    extend (ModMod (Scope inner_scope inner_mods)) =+      -- XXX: perhaps hacky fix for #1653.  We need to impose the+      -- substitutions of abstract types from outside, because the+      -- inner module may have some incorrect substitutions in some+      -- cases.  Our treatment of abstract types is completely whack+      -- and should be fixed.+      ModMod $ Scope (scopeSubsts scope <> inner_scope) inner_mods+    extend m = m  lookupMod :: QualName VName -> TransformM Mod lookupMod mname = either error pure . lookupMod' mname =<< askScope

@@ -1578,6 +1578,7 @@           handleOps,           handleSOACs,           handleAccs,+          handleAD,           handleRest         ]   msum [h args $ baseString $ qualLeaf qname | h <- handlers]@@ -1745,6 +1746,19 @@       vs <- internaliseExp "acc_v" v       fmap pure $ letSubExp desc $ BasicOp $ UpdateAcc acc' [i'] vs     handleAccs _ _ = Nothing++    handleAD [TupLit [f, x, v] _] fname+      | fname `elem` ["jvp2", "vjp2"] = Just $ \desc -> do+          x' <- internaliseExp "ad_x" x+          v' <- internaliseExp "ad_v" v+          xts <- mapM subExpType x'+          (ps, body, ret) <- internaliseLambda f xts+          let lam = I.Lambda ps body ret+          fmap (map I.Var) . letTupExp desc . Op $+            case fname of+              "jvp2" -> JVP lam x' v'+              _ -> VJP lam x' v'+    handleAD _ _ = Nothing      handleRest [E.TupLit [a, si, v] _] "scatter" = Just $ scatterF 1 a si v     handleRest [E.TupLit [a, si, v] _] "scatter_2d" = Just $ scatterF 2 a si v

@@ -2,61 +2,88 @@ -- the Futhark program managed by the language server.  The challenge -- here is that if the program becomes type-invalid, we want to keep -- the old state around.-module Futhark.LSP.Compile (tryTakeStateFromMVar, tryReCompile) where+module Futhark.LSP.Compile (tryTakeStateFromIORef, tryReCompile) where -import Control.Concurrent.MVar (MVar, putMVar, takeMVar) import Control.Monad.IO.Class (MonadIO (liftIO))+import Data.IORef (IORef, readIORef, writeIORef) import qualified Data.Map as M+import Data.Maybe (fromMaybe) import qualified Data.Text as T-import Futhark.Compiler.Program (LoadedProg, lpWarnings, noLoadedProg, reloadProg)+import Futhark.Compiler.Program (LoadedProg, lpFilePaths, lpWarnings, noLoadedProg, reloadProg) import Futhark.LSP.Diagnostic (diagnosticSource, maxDiagnostic, publishErrorDiagnostics, publishWarningDiagnostics)-import Futhark.LSP.State (State (..), emptyState)+import Futhark.LSP.State (State (..), emptyState, updateStaleContent, updateStaleMapping)+import Futhark.LSP.Tool (computeMapping) import Futhark.Util (debug) import Language.Futhark.Warnings (listWarnings)-import Language.LSP.Server (LspT, flushDiagnosticsBySource, getVirtualFiles)-import Language.LSP.Types (fromNormalizedFilePath, uriToNormalizedFilePath)+import Language.LSP.Server (LspT, flushDiagnosticsBySource, getVirtualFile, getVirtualFiles)+import Language.LSP.Types+  ( filePathToUri,+    fromNormalizedFilePath,+    toNormalizedUri,+    uriToNormalizedFilePath,+  ) import Language.LSP.VFS (VFS (vfsMap), virtualFileText) --- | Try to take state from MVar, if it's empty, try to compile.-tryTakeStateFromMVar :: MVar State -> Maybe FilePath -> LspT () IO State-tryTakeStateFromMVar state_mvar file_path = do-  old_state <- liftIO $ takeMVar state_mvar+-- | Try to take state from IORef, if it's empty, try to compile.+tryTakeStateFromIORef :: IORef State -> Maybe FilePath -> LspT () IO State+tryTakeStateFromIORef state_mvar file_path = do+  old_state <- liftIO $ readIORef state_mvar   case stateProgram old_state of     Nothing -> do-      new_state <- tryCompile file_path (State $ Just noLoadedProg)-      liftIO $ putMVar state_mvar new_state+      new_state <- tryCompile old_state file_path noLoadedProg+      liftIO $ writeIORef state_mvar new_state       pure new_state-    Just _ -> do-      liftIO $ putMVar state_mvar old_state-      pure old_state+    Just prog -> do+      -- If this is in the context of some file that is not part of+      -- the program, try to reload the program from that file.+      let files = lpFilePaths prog+      state <- case file_path of+        Just file_path'+          | file_path' `notElem` files -> do+              debug $ "File not part of program: " <> show file_path'+              debug $ "Program contains: " <> show files+              tryCompile old_state file_path noLoadedProg+        _ -> pure old_state+      liftIO $ writeIORef state_mvar state+      pure state  -- | Try to (re)-compile, replace old state if successful.-tryReCompile :: MVar State -> Maybe FilePath -> LspT () IO ()+tryReCompile :: IORef State -> Maybe FilePath -> LspT () IO () tryReCompile state_mvar file_path = do   debug "(Re)-compiling ..."-  old_state <- liftIO $ takeMVar state_mvar-  new_state <- tryCompile file_path old_state+  old_state <- liftIO $ readIORef state_mvar+  let loaded_prog = getLoadedProg old_state+  new_state <- tryCompile old_state file_path loaded_prog   case stateProgram new_state of     Nothing -> do       debug "Failed to (re)-compile, using old state or Nothing"-      liftIO $ putMVar state_mvar old_state+      debug $ "Computing PositionMapping for: " <> show file_path+      mapping <- computeMapping old_state file_path+      liftIO $ writeIORef state_mvar $ updateStaleMapping file_path mapping old_state     Just _ -> do       debug "(Re)-compile successful"-      liftIO $ putMVar state_mvar new_state+      liftIO $ writeIORef state_mvar new_state  -- | Try to compile, publish diagnostics on warnings and errors, return newly compiled state. --  Single point where the compilation is done, and shouldn't be exported.-tryCompile :: Maybe FilePath -> State -> LspT () IO State-tryCompile Nothing _ = pure emptyState-tryCompile (Just path) state = do-  let old_loaded_prog = getLoadedProg state+tryCompile :: State -> Maybe FilePath -> LoadedProg -> LspT () IO State+tryCompile _ Nothing _ = pure emptyState+tryCompile state (Just path) old_loaded_prog = do+  debug $ "Reloading program from " <> show path   vfs <- getVirtualFiles-  res <- liftIO $ reloadProg old_loaded_prog [path] (transformVFS vfs)+  res <- liftIO $ reloadProg old_loaded_prog [path] (transformVFS vfs) -- NOTE: vfs only keeps track of current opened files   flushDiagnosticsBySource maxDiagnostic diagnosticSource   case res of     Right new_loaded_prog -> do       publishWarningDiagnostics $ listWarnings $ lpWarnings new_loaded_prog-      pure $ State (Just new_loaded_prog)+      maybe_virtual_file <- getVirtualFile $ toNormalizedUri $ filePathToUri path+      case maybe_virtual_file of+        Nothing -> pure $ State (Just new_loaded_prog) (staleData state) -- should never happen+        Just virtual_file ->+          pure $ updateStaleContent path virtual_file new_loaded_prog state+    -- Preserve files that have been opened should be enoguth.+    -- But still might need an update on re-compile logic, don't discard all state afterwards,+    -- try to compile from root file, if there is a depencency relatetion, improve performance and provide more dignostic.     Left prog_error -> do       debug "Compilation failed, publishing diagnostics"       publishErrorDiagnostics prog_error@@ -77,5 +104,4 @@     (vfsMap vfs)  getLoadedProg :: State -> LoadedProg-getLoadedProg (State (Just loaded_prog)) = loaded_prog-getLoadedProg (State Nothing) = noLoadedProg+getLoadedProg state = fromMaybe noLoadedProg (stateProgram state)

@@ -3,11 +3,11 @@ -- | The handlers exposed by the language server. module Futhark.LSP.Handlers (handlers) where -import Control.Concurrent.MVar (MVar) import Control.Lens ((^.)) import Data.Aeson.Types (Value (Array, String))+import Data.IORef import qualified Data.Vector as V-import Futhark.LSP.Compile (tryReCompile, tryTakeStateFromMVar)+import Futhark.LSP.Compile (tryReCompile, tryTakeStateFromIORef) import Futhark.LSP.State (State (..)) import Futhark.LSP.Tool (findDefinitionRange, getHoverInfoFromState) import Futhark.Util (debug)@@ -15,36 +15,35 @@ import Language.LSP.Types import Language.LSP.Types.Lens (HasUri (uri)) --- | Given an 'MVar' tracking the state, produce a set of handlers.+-- | Given an 'IORef' tracking the state, produce a set of handlers. -- When we want to add more features to the language server, this is -- the thing to change.-handlers :: MVar State -> Handlers (LspM ())+handlers :: IORef State -> Handlers (LspM ()) handlers state_mvar =   mconcat     [ onInitializeHandler,-      onHoverHandler state_mvar,       onDocumentOpenHandler state_mvar,       onDocumentCloseHandler,       onDocumentSaveHandler state_mvar,       onDocumentChangeHandler state_mvar,       onDocumentFocusHandler state_mvar,       goToDefinitionHandler state_mvar,-      onDocumentChangeHandler state_mvar+      onHoverHandler state_mvar     ]  onInitializeHandler :: Handlers (LspM ()) onInitializeHandler = notificationHandler SInitialized $ \_msg -> debug "Initialized" -onHoverHandler :: MVar State -> Handlers (LspM ())+onHoverHandler :: IORef State -> Handlers (LspM ()) onHoverHandler state_mvar = requestHandler STextDocumentHover $ \req responder -> do-  debug "Got hover request"   let RequestMessage _ _ _ (HoverParams doc pos _workDone) = req       Position l c = pos       file_path = uriToFilePath $ doc ^. uri-  state <- tryTakeStateFromMVar state_mvar file_path+  debug $ "Got hover request: " <> show (file_path, pos)+  state <- tryTakeStateFromIORef state_mvar file_path   responder $ Right $ getHoverInfoFromState state file_path (fromEnum l + 1) (fromEnum c + 1) -onDocumentFocusHandler :: MVar State -> Handlers (LspM ())+onDocumentFocusHandler :: IORef State -> Handlers (LspM ()) onDocumentFocusHandler state_mvar = notificationHandler (SCustomMethod "custom/onFocusTextDocument") $ \msg -> do   debug "Got custom request: onFocusTextDocument"   let NotificationMessage _ _ (Array vector_param) = msg@@ -52,33 +51,33 @@   debug $ show focused_uri   tryReCompile state_mvar (uriToFilePath (Uri focused_uri)) -goToDefinitionHandler :: MVar State -> Handlers (LspM ())+goToDefinitionHandler :: IORef State -> Handlers (LspM ()) goToDefinitionHandler state_mvar = requestHandler STextDocumentDefinition $ \req responder -> do-  debug "Got goto definition request"   let RequestMessage _ _ _ (DefinitionParams doc pos _workDone _partial) = req       Position l c = pos       file_path = uriToFilePath $ doc ^. uri-  state <- tryTakeStateFromMVar state_mvar file_path+  debug $ "Got goto definition: " <> show (file_path, pos)+  state <- tryTakeStateFromIORef state_mvar file_path   case findDefinitionRange state file_path (fromEnum l + 1) (fromEnum c + 1) of     Nothing -> responder $ Right $ InR $ InL $ List []     Just loc -> do       debug $ show loc       responder $ Right $ InL loc -onDocumentSaveHandler :: MVar State -> Handlers (LspM ())+onDocumentSaveHandler :: IORef State -> Handlers (LspM ()) onDocumentSaveHandler state_mvar = notificationHandler STextDocumentDidSave $ \msg -> do   let NotificationMessage _ _ (DidSaveTextDocumentParams doc _text) = msg       file_path = uriToFilePath $ doc ^. uri   debug $ "Saved document: " ++ show doc   tryReCompile state_mvar file_path -onDocumentChangeHandler :: MVar State -> Handlers (LspM ())+onDocumentChangeHandler :: IORef State -> Handlers (LspM ()) onDocumentChangeHandler state_mvar = notificationHandler STextDocumentDidChange $ \msg -> do   let NotificationMessage _ _ (DidChangeTextDocumentParams doc _content) = msg       file_path = uriToFilePath $ doc ^. uri   tryReCompile state_mvar file_path -onDocumentOpenHandler :: MVar State -> Handlers (LspM ())+onDocumentOpenHandler :: IORef State -> Handlers (LspM ()) onDocumentOpenHandler state_mvar = notificationHandler STextDocumentDidOpen $ \msg -> do   let NotificationMessage _ _ (DidOpenTextDocumentParams doc) = msg       file_path = uriToFilePath $ doc ^. uri

@@ -0,0 +1,80 @@+-- | Provide mapping between position in stale content and current.+module Futhark.LSP.PositionMapping+  ( mappingFromDiff,+    PositionMapping,+    toStalePos,+    toCurrentLoc,+    StaleFile (..),+  )+where++import Data.Algorithm.Diff (Diff, PolyDiff (Both, First, Second), getDiff)+import Data.Bifunctor (Bifunctor (bimap, first, second))+import qualified Data.Text as T+import qualified Data.Vector as V+import Futhark.Util.Loc (Loc (Loc), Pos (Pos))+import Language.LSP.VFS (VirtualFile)++-- | A mapping between current file content and the stale (last successful compiled) file content.+-- Currently, only supports entire line mapping,+-- more detailed mapping might be achieved via referring to haskell-language-server@efb4b94+data PositionMapping = PositionMapping+  { -- | The mapping from stale position to current.+    -- e.g. staleToCurrent[2] = 4 means "line 2" in the stale file, corresponds to "line 4" in the current file.+    staleToCurrent :: V.Vector Int,+    -- | The mapping from current position to stale.+    currentToStale :: V.Vector Int+  }+  deriving (Show)++-- | Stale text document stored in state.+data StaleFile = StaleFile+  { -- | The last successfully compiled file content.+    -- Using VirtualFile for convenience, we can use anything with {version, content}+    staleContent :: VirtualFile,+    -- | PositionMapping between current and stale file content.+    -- Nothing if last type-check is successful.+    staleMapping :: Maybe PositionMapping+  }+  deriving (Show)++-- | Compute PositionMapping using the diff between two texts.+mappingFromDiff :: [T.Text] -> [T.Text] -> PositionMapping+mappingFromDiff stale current = do+  let (stale_to_current, current_to_stale) = rawMapping (getDiff stale current) 0 0+  PositionMapping (V.fromList stale_to_current) (V.fromList current_to_stale)+  where+    rawMapping :: [Diff T.Text] -> Int -> Int -> ([Int], [Int])+    rawMapping [] _ _ = ([], [])+    rawMapping (Both _ _ : xs) lold lnew = bimap (lnew :) (lold :) $ rawMapping xs (lold + 1) (lnew + 1)+    rawMapping (First _ : xs) lold lnew = first (-1 :) $ rawMapping xs (lold + 1) lnew+    rawMapping (Second _ : xs) lold lnew = second (-1 :) $ rawMapping xs lold (lnew + 1)++-- | Transform current Pos to the stale pos for query+-- Note: line and col in Pos is larger by one+toStalePos :: Maybe PositionMapping -> Pos -> Maybe Pos+toStalePos (Just (PositionMapping _ current_to_stale)) pos =+  if l > Prelude.length current_to_stale+    then Nothing+    else Just $ Pos file (V.unsafeIndex current_to_stale (l - 1) + 1) c o+  where+    Pos file l c o = pos+toStalePos Nothing pos = Just pos++-- some refactoring might be needed, same logic as toStalePos+toCurrentPos :: Maybe PositionMapping -> Pos -> Maybe Pos+toCurrentPos (Just (PositionMapping stale_to_current _)) pos =+  if l > Prelude.length stale_to_current+    then Nothing+    else Just $ Pos file (V.unsafeIndex stale_to_current (l - 1) + 1) c o+  where+    Pos file l c o = pos+toCurrentPos Nothing pos = Just pos++-- | Transform stale Loc gotten from stale AST to current Loc.+toCurrentLoc :: Maybe PositionMapping -> Loc -> Maybe Loc+toCurrentLoc mapping loc = do+  let Loc start end = loc+  current_start <- toCurrentPos mapping start+  current_end <- toCurrentPos mapping end+  Just $ Loc current_start current_end

@@ -2,17 +2,53 @@ module Futhark.LSP.State   ( State (..),     emptyState,+    getStaleContent,+    getStaleMapping,+    updateStaleContent,+    updateStaleMapping,   ) where +import qualified Data.Map as M import Futhark.Compiler.Program (LoadedProg)+import Futhark.LSP.PositionMapping (PositionMapping, StaleFile (..))+import Language.LSP.VFS (VirtualFile)  -- | The state of the language server.-newtype State = State+data State = State   { -- | The loaded program.-    stateProgram :: Maybe LoadedProg+    stateProgram :: Maybe LoadedProg,+    -- | The stale data, stored to provide PositionMapping when requested.+    -- All files that have been opened have an entry.+    staleData :: M.Map FilePath StaleFile   }  -- | Initial state. emptyState :: State-emptyState = State Nothing+emptyState = State Nothing M.empty++-- | Get the contents of a stale (last successfully complied) file's contents.+getStaleContent :: State -> FilePath -> Maybe VirtualFile+getStaleContent state file_path = (Just . staleContent) =<< M.lookup file_path (staleData state)++-- | Get the PositionMapping for a file.+getStaleMapping :: State -> FilePath -> Maybe PositionMapping+getStaleMapping state file_path = staleMapping =<< M.lookup file_path (staleData state)++-- | Update the state with another pair of file_path and contents.+-- Could do a clean up becausae there is no need to store files that are not in lpFilePaths prog.+updateStaleContent :: FilePath -> VirtualFile -> LoadedProg -> State -> State+updateStaleContent file_path file_content loadedProg state =+  -- NOTE: insert will replace the old value if the key already exists.+  -- updateStaleContent is only called after a successful type-check,+  -- so the PositionsMapping should be Nothing here, it's calculated after failed type-check.+  State (Just loadedProg) (M.insert file_path (StaleFile file_content Nothing) (staleData state))++-- | Update the state with another pair of file_path and PositionMapping.+updateStaleMapping :: Maybe FilePath -> Maybe PositionMapping -> State -> State+updateStaleMapping (Just file_path) mapping state = do+  case M.lookup file_path (staleData state) of+    Nothing -> state -- Only happends when the file have never been successfully type-checked before.+    Just (StaleFile file_content _mapping) ->+      State (stateProgram state) (M.insert file_path (StaleFile file_content mapping) (staleData state))+updateStaleMapping _ _ state = state

@@ -8,11 +8,19 @@     rangeFromSrcLoc,     rangeFromLoc,     posToUri,+    computeMapping,   ) where +import qualified Data.Text as T import Futhark.Compiler.Program (lpImports)-import Futhark.LSP.State (State (..))+import Futhark.LSP.PositionMapping+  ( PositionMapping,+    mappingFromDiff,+    toCurrentLoc,+    toStalePos,+  )+import Futhark.LSP.State (State (..), getStaleContent, getStaleMapping) import Futhark.Util.Loc (Loc (Loc, NoLoc), Pos (Pos), SrcLoc, locOf) import Futhark.Util.Pretty (prettyText) import Language.Futhark.Prop (isBuiltinLoc)@@ -22,7 +30,9 @@     atPos,     boundLoc,   )+import Language.LSP.Server (LspM, getVirtualFile) import Language.LSP.Types+import Language.LSP.VFS (VirtualFile, virtualFileText, virtualFileVersion)  -- | Retrieve hover info for the definition referenced at the given -- file at the given line and column number (the two 'Int's).@@ -53,11 +63,53 @@     else Just $ Location (filePathToUri file_path) (rangeFromLoc loc) findDefinitionRange _ _ _ _ = Nothing +-- | Query the AST for information at certain Pos. queryAtPos :: State -> Pos -> Maybe AtPos-queryAtPos state pos =-  case stateProgram state of-    Nothing -> Nothing-    Just loaded_prog -> atPos (lpImports loaded_prog) pos+queryAtPos state pos = do+  let Pos path _ _ _ = pos+      mapping = getStaleMapping state path+  loaded_prog <- stateProgram state+  stale_pos <- toStalePos mapping pos+  query_result <- atPos (lpImports loaded_prog) stale_pos+  updateAtPos mapping query_result+  where+    -- Update the 'AtPos' with the current mapping.+    updateAtPos :: Maybe PositionMapping -> AtPos -> Maybe AtPos+    updateAtPos mapping (AtName qn (Just def) loc) = do+      let def_loc = boundLoc def+          Loc (Pos def_file _ _ _) _ = def_loc+          Pos current_file _ _ _ = pos+      current_loc <- toCurrentLoc mapping loc+      if def_file == current_file+        then do+          current_def_loc <- toCurrentLoc mapping def_loc+          Just $ AtName qn (Just (updateBoundLoc def current_def_loc)) current_loc+        else do+          -- Defined in another file, get the corresponding PositionMapping.+          let def_mapping = getStaleMapping state def_file+          current_def_loc <- toCurrentLoc def_mapping def_loc+          Just $ AtName qn (Just (updateBoundLoc def current_def_loc)) current_loc+    updateAtPos _ _ = Nothing++    updateBoundLoc :: BoundTo -> Loc -> BoundTo+    updateBoundLoc (BoundTerm t _loc) current_loc = BoundTerm t current_loc+    updateBoundLoc (BoundModule _loc) current_loc = BoundModule current_loc+    updateBoundLoc (BoundModuleType _loc) current_loc = BoundModuleType current_loc+    updateBoundLoc (BoundType _loc) current_loc = BoundType current_loc++-- | Entry point for computing PositionMapping.+computeMapping :: State -> Maybe FilePath -> LspM () (Maybe PositionMapping)+computeMapping state (Just file_path) = do+  virtual_file <- getVirtualFile $ toNormalizedUri $ filePathToUri file_path+  pure $ getMapping (getStaleContent state file_path) virtual_file+  where+    getMapping :: Maybe VirtualFile -> Maybe VirtualFile -> Maybe PositionMapping+    getMapping (Just stale_file) (Just current_file) =+      if virtualFileVersion stale_file == virtualFileVersion current_file+        then Nothing -- Happens when other files (e.g. dependencies) fail to type-check.+        else Just $ mappingFromDiff (T.lines $ virtualFileText stale_file) (T.lines $ virtualFileText current_file)+    getMapping _ _ = Nothing+computeMapping _ _ = pure Nothing  -- | Convert a Futhark 'Pos' to an LSP 'Uri'. posToUri :: Pos -> Uri