diff --git a/docs/conf.py b/docs/conf.py --- a/docs/conf.py +++ b/docs/conf.py @@ -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() diff --git a/docs/error-index.rst b/docs/error-index.rst --- a/docs/error-index.rst +++ b/docs/error-index.rst @@ -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 `, 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 `, 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 diff --git a/docs/installation.rst b/docs/installation.rst --- a/docs/installation.rst +++ b/docs/installation.rst @@ -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 `_ 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 -`_. 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 -------------------------------------- diff --git a/docs/man/futhark-test.rst b/docs/man/futhark-test.rst --- a/docs/man/futhark-test.rst +++ b/docs/man/futhark-test.rst @@ -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. diff --git a/futhark.cabal b/futhark.cabal --- a/futhark.cabal +++ b/futhark.cabal @@ -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 diff --git a/prelude/ad.fut b/prelude/ad.fut new file mode 100644 --- /dev/null +++ b/prelude/ad.fut @@ -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 diff --git a/prelude/array.fut b/prelude/array.fut --- a/prelude/array.fut +++ b/prelude/array.fut @@ -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. diff --git a/prelude/prelude.fut b/prelude/prelude.fut --- a/prelude/prelude.fut +++ b/prelude/prelude.fut @@ -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 diff --git a/prelude/soacs.fut b/prelude/soacs.fut --- a/prelude/soacs.fut +++ b/prelude/soacs.fut @@ -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) diff --git a/prelude/zip.fut b/prelude/zip.fut --- a/prelude/zip.fut +++ b/prelude/zip.fut @@ -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) = diff --git a/rts/c/chaselev.h b/rts/c/chaselev.h deleted file mode 100644 --- a/rts/c/chaselev.h +++ /dev/null @@ -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 -#include -#include - - -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 diff --git a/rts/c/cuda.h b/rts/c/cuda.h --- a/rts/c/cuda.h +++ b/rts/c/cuda.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); diff --git a/rts/python/opencl.py b/rts/python/opencl.py --- a/rts/python/opencl.py +++ b/rts/python/opencl.py @@ -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'] diff --git a/src/Futhark/AD/Derivatives.hs b/src/Futhark/AD/Derivatives.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/AD/Derivatives.hs @@ -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 diff --git a/src/Futhark/AD/Fwd.hs b/src/Futhark/AD/Fwd.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/AD/Fwd.hs @@ -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) diff --git a/src/Futhark/AD/Rev.hs b/src/Futhark/AD/Rev.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/AD/Rev.hs @@ -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. diff --git a/src/Futhark/AD/Rev/Loop.hs b/src/Futhark/AD/Rev/Loop.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/AD/Rev/Loop.hs @@ -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 diff --git a/src/Futhark/AD/Rev/Map.hs b/src/Futhark/AD/Rev/Map.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/AD/Rev/Map.hs @@ -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 diff --git a/src/Futhark/AD/Rev/Monad.hs b/src/Futhark/AD/Rev/Monad.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/AD/Rev/Monad.hs @@ -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. diff --git a/src/Futhark/AD/Rev/Reduce.hs b/src/Futhark/AD/Rev/Reduce.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/AD/Rev/Reduce.hs @@ -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) diff --git a/src/Futhark/AD/Rev/SOAC.hs b/src/Futhark/AD/Rev/SOAC.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/AD/Rev/SOAC.hs @@ -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 diff --git a/src/Futhark/AD/Rev/Scan.hs b/src/Futhark/AD/Rev/Scan.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/AD/Rev/Scan.hs @@ -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 diff --git a/src/Futhark/AD/Rev/Scatter.hs b/src/Futhark/AD/Rev/Scatter.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/AD/Rev/Scatter.hs @@ -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 diff --git a/src/Futhark/Analysis/LastUse.hs b/src/Futhark/Analysis/LastUse.hs --- a/src/Futhark/Analysis/LastUse.hs +++ b/src/Futhark/Analysis/LastUse.hs @@ -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) => diff --git a/src/Futhark/Analysis/Metrics.hs b/src/Futhark/Analysis/Metrics.hs --- a/src/Futhark/Analysis/Metrics.hs +++ b/src/Futhark/Analysis/Metrics.hs @@ -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 diff --git a/src/Futhark/Analysis/PrimExp.hs b/src/Futhark/Analysis/PrimExp.hs --- a/src/Futhark/Analysis/PrimExp.hs +++ b/src/Futhark/Analysis/PrimExp.hs @@ -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 ~+~, ~-~ diff --git a/src/Futhark/Bench.hs b/src/Futhark/Bench.hs --- a/src/Futhark/Bench.hs +++ b/src/Futhark/Bench.hs @@ -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 diff --git a/src/Futhark/CLI/Bench.hs b/src/Futhark/CLI/Bench.hs --- a/src/Futhark/CLI/Bench.hs +++ b/src/Futhark/CLI/Bench.hs @@ -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. diff --git a/src/Futhark/CLI/Dev.hs b/src/Futhark/CLI/Dev.hs --- a/src/Futhark/CLI/Dev.hs +++ b/src/Futhark/CLI/Dev.hs @@ -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 [], diff --git a/src/Futhark/CLI/LSP.hs b/src/Futhark/CLI/LSP.hs --- a/src/Futhark/CLI/LSP.hs +++ b/src/Futhark/CLI/LSP.hs @@ -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 = diff --git a/src/Futhark/CLI/Literate.hs b/src/Futhark/CLI/Literate.hs --- a/src/Futhark/CLI/Literate.hs +++ b/src/Futhark/CLI/Literate.hs @@ -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) diff --git a/src/Futhark/CLI/Main.hs b/src/Futhark/CLI/Main.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CLI/Main.hs @@ -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 $ + [" 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 diff --git a/src/Futhark/CLI/Test.hs b/src/Futhark/CLI/Test.hs --- a/src/Futhark/CLI/Test.hs +++ b/src/Futhark/CLI/Test.hs @@ -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 = diff --git a/src/Futhark/CodeGen/Backends/CCUDA.hs b/src/Futhark/CodeGen/Backends/CCUDA.hs --- a/src/Futhark/CodeGen/Backends/CCUDA.hs +++ b/src/Futhark/CodeGen/Backends/CCUDA.hs @@ -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." diff --git a/src/Futhark/CodeGen/Backends/COpenCL.hs b/src/Futhark/CodeGen/Backends/COpenCL.hs --- a/src/Futhark/CodeGen/Backends/COpenCL.hs +++ b/src/Futhark/CodeGen/Backends/COpenCL.hs @@ -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." diff --git a/src/Futhark/CodeGen/Backends/GenericPython.hs b/src/Futhark/CodeGen/Backends/GenericPython.hs --- a/src/Futhark/CodeGen/Backends/GenericPython.hs +++ b/src/Futhark/CodeGen/Backends/GenericPython.hs @@ -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) diff --git a/src/Futhark/CodeGen/ImpCode.hs b/src/Futhark/CodeGen/ImpCode.hs --- a/src/Futhark/CodeGen/ImpCode.hs +++ b/src/Futhark/CodeGen/ImpCode.hs @@ -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 diff --git a/src/Futhark/CodeGen/ImpGen/GPU.hs b/src/Futhark/CodeGen/ImpGen/GPU.hs --- a/src/Futhark/CodeGen/ImpGen/GPU.hs +++ b/src/Futhark/CodeGen/ImpGen/GPU.hs @@ -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 " diff --git a/src/Futhark/CodeGen/OpenCL/Heuristics.hs b/src/Futhark/CodeGen/OpenCL/Heuristics.hs --- a/src/Futhark/CodeGen/OpenCL/Heuristics.hs +++ b/src/Futhark/CodeGen/OpenCL/Heuristics.hs @@ -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, diff --git a/src/Futhark/CodeGen/RTS/C.hs b/src/Futhark/CodeGen/RTS/C.hs --- a/src/Futhark/CodeGen/RTS/C.hs +++ b/src/Futhark/CodeGen/RTS/C.hs @@ -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 diff --git a/src/Futhark/Compiler.hs b/src/Futhark/Compiler.hs --- a/src/Futhark/Compiler.hs +++ b/src/Futhark/Compiler.hs @@ -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. diff --git a/src/Futhark/Compiler/CLI.hs b/src/Futhark/Compiler/CLI.hs --- a/src/Futhark/Compiler/CLI.hs +++ b/src/Futhark/Compiler/CLI.hs @@ -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 = diff --git a/src/Futhark/Compiler/Config.hs b/src/Futhark/Compiler/Config.hs --- a/src/Futhark/Compiler/Config.hs +++ b/src/Futhark/Compiler/Config.hs @@ -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 diff --git a/src/Futhark/Compiler/Program.hs b/src/Futhark/Compiler/Program.hs --- a/src/Futhark/Compiler/Program.hs +++ b/src/Futhark/Compiler/Program.hs @@ -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. diff --git a/src/Futhark/Construct.hs b/src/Futhark/Construct.hs --- a/src/Futhark/Construct.hs +++ b/src/Futhark/Construct.hs @@ -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 :: diff --git a/src/Futhark/Doc/Generator.hs b/src/Futhark/Doc/Generator.hs --- a/src/Futhark/Doc/Generator.hs +++ b/src/Futhark/Doc/Generator.hs @@ -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 diff --git a/src/Futhark/IR/Aliases.hs b/src/Futhark/IR/Aliases.hs --- a/src/Futhark/IR/Aliases.hs +++ b/src/Futhark/IR/Aliases.hs @@ -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 diff --git a/src/Futhark/IR/GPU/Op.hs b/src/Futhark/IR/GPU/Op.hs --- a/src/Futhark/IR/GPU/Op.hs +++ b/src/Futhark/IR/GPU/Op.hs @@ -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 + ] diff --git a/src/Futhark/IR/GPU/Simplify.hs b/src/Futhark/IR/GPU/Simplify.hs --- a/src/Futhark/IR/GPU/Simplify.hs +++ b/src/Futhark/IR/GPU/Simplify.hs @@ -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 diff --git a/src/Futhark/IR/Mem/Simplify.hs b/src/Futhark/IR/Mem/Simplify.hs --- a/src/Futhark/IR/Mem/Simplify.hs +++ b/src/Futhark/IR/Mem/Simplify.hs @@ -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 diff --git a/src/Futhark/IR/Parse.hs b/src/Futhark/IR/Parse.hs --- a/src/Futhark/IR/Parse.hs +++ b/src/Futhark/IR/Parse.hs @@ -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) diff --git a/src/Futhark/IR/Primitive.hs b/src/Futhark/IR/Primitive.hs --- a/src/Futhark/IR/Primitive.hs +++ b/src/Futhark/IR/Primitive.hs @@ -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 diff --git a/src/Futhark/IR/Prop/Aliases.hs b/src/Futhark/IR/Prop/Aliases.hs --- a/src/Futhark/IR/Prop/Aliases.hs +++ b/src/Futhark/IR/Prop/Aliases.hs @@ -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 diff --git a/src/Futhark/IR/Prop/Names.hs b/src/Futhark/IR/Prop/Names.hs --- a/src/Futhark/IR/Prop/Names.hs +++ b/src/Futhark/IR/Prop/Names.hs @@ -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 diff --git a/src/Futhark/IR/Prop/Types.hs b/src/Futhark/IR/Prop/Types.hs --- a/src/Futhark/IR/Prop/Types.hs +++ b/src/Futhark/IR/Prop/Types.hs @@ -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. diff --git a/src/Futhark/IR/SOACS.hs b/src/Futhark/IR/SOACS.hs --- a/src/Futhark/IR/SOACS.hs +++ b/src/Futhark/IR/SOACS.hs @@ -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 diff --git a/src/Futhark/IR/SOACS/SOAC.hs b/src/Futhark/IR/SOACS/SOAC.hs --- a/src/Futhark/IR/SOACS/SOAC.hs +++ b/src/Futhark/IR/SOACS/SOAC.hs @@ -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 -> diff --git a/src/Futhark/IR/SOACS/Simplify.hs b/src/Futhark/IR/SOACS/Simplify.hs --- a/src/Futhark/IR/SOACS/Simplify.hs +++ b/src/Futhark/IR/SOACS/Simplify.hs @@ -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) diff --git a/src/Futhark/IR/SegOp.hs b/src/Futhark/IR/SegOp.hs --- a/src/Futhark/IR/SegOp.hs +++ b/src/Futhark/IR/SegOp.hs @@ -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} diff --git a/src/Futhark/IR/Syntax.hs b/src/Futhark/IR/Syntax.hs --- a/src/Futhark/IR/Syntax.hs +++ b/src/Futhark/IR/Syntax.hs @@ -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 diff --git a/src/Futhark/IR/Syntax/Core.hs b/src/Futhark/IR/Syntax/Core.hs --- a/src/Futhark/IR/Syntax/Core.hs +++ b/src/Futhark/IR/Syntax/Core.hs @@ -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 diff --git a/src/Futhark/Internalise/Defunctorise.hs b/src/Futhark/Internalise/Defunctorise.hs --- a/src/Futhark/Internalise/Defunctorise.hs +++ b/src/Futhark/Internalise/Defunctorise.hs @@ -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 diff --git a/src/Futhark/Internalise/Exps.hs b/src/Futhark/Internalise/Exps.hs --- a/src/Futhark/Internalise/Exps.hs +++ b/src/Futhark/Internalise/Exps.hs @@ -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 diff --git a/src/Futhark/LSP/Compile.hs b/src/Futhark/LSP/Compile.hs --- a/src/Futhark/LSP/Compile.hs +++ b/src/Futhark/LSP/Compile.hs @@ -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) diff --git a/src/Futhark/LSP/Handlers.hs b/src/Futhark/LSP/Handlers.hs --- a/src/Futhark/LSP/Handlers.hs +++ b/src/Futhark/LSP/Handlers.hs @@ -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 diff --git a/src/Futhark/LSP/PositionMapping.hs b/src/Futhark/LSP/PositionMapping.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/LSP/PositionMapping.hs @@ -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 diff --git a/src/Futhark/LSP/State.hs b/src/Futhark/LSP/State.hs --- a/src/Futhark/LSP/State.hs +++ b/src/Futhark/LSP/State.hs @@ -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 diff --git a/src/Futhark/LSP/Tool.hs b/src/Futhark/LSP/Tool.hs --- a/src/Futhark/LSP/Tool.hs +++ b/src/Futhark/LSP/Tool.hs @@ -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 diff --git a/src/Futhark/Optimise/BlkRegTiling.hs b/src/Futhark/Optimise/BlkRegTiling.hs --- a/src/Futhark/Optimise/BlkRegTiling.hs +++ b/src/Futhark/Optimise/BlkRegTiling.hs @@ -21,909 +21,1205 @@ import Control.Monad.Reader import qualified Data.List as L -import qualified Data.Map.Strict as M -import Data.Maybe -import qualified Data.Sequence as Seq -import Futhark.IR.GPU -import Futhark.MonadFreshNames -import Futhark.Optimise.TileLoops.Shared -import Futhark.Tools -import Futhark.Transform.Rename - -mmBlkRegTiling :: Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU)) -mmBlkRegTiling (Let pat aux (Op (SegOp (SegMap SegThread {} seg_space ts old_kbody)))) - | KernelBody () kstms [Returns ResultMaySimplify cs (Var res_nm)] <- old_kbody, - cs == mempty, - -- check kernel has one result of primitive type - [res_tp] <- ts, - primType res_tp, - -- build the variance table, that records, for - -- each variable name, the variables it depends on - initial_variance <- M.map mempty $ scopeOfSegSpace seg_space, - variance <- varianceInStms initial_variance kstms, - -- check that the code fits the pattern having: - -- some `code1`, followed by one Screma SOAC, followed by some `code2` - (code1, Just screma_stmt, code2) <- matchCodeStreamCode kstms, - Let pat_redomap _ (Op _) <- screma_stmt, - -- checks that the Screma SOAC is actually a redomap and normalizes it - Just (common_dim, arrs, (_, red_lam, red_nes, map_lam)) <- isTileableRedomap screma_stmt, - -- check that exactly two 1D arrays are streamed thorugh redomap, - -- and the result of redomap is one scalar - -- !!!I need to rearrange this whole thing!!! including inp_A and inp_B - length arrs == 2 && length red_nes == 1, - [map_t1t, map_t2t] <- map paramDec $ lambdaParams map_lam, - [red_t1, _] <- map paramDec $ lambdaParams red_lam, - primType map_t1t && primType map_t2t && primType red_t1, - map_t1 <- elemType map_t1t, - map_t2 <- elemType map_t2t, - -- checks that the input arrays to redomap are variant to - -- exactly one of the two innermost dimensions of the kernel - Just var_dims <- isInvarTo1of2InnerDims mempty seg_space variance arrs, - -- get the variables on which the first result of redomap depends on - [redomap_orig_res] <- patElems pat_redomap, - Just res_red_var <- M.lookup (patElemName redomap_orig_res) variance, -- variance of the reduce result - - -- we furthermore check that code1 is only formed by - -- 1. statements that slice some globally-declared arrays - -- to produce the input for the redomap, and - -- 2. potentially some statements on which the redomap - -- is independent; these are recorded in `code2''` - Just (code2'', tab_inv_stm) <- - foldl - (processIndirections (namesFromList arrs) res_red_var) - (Just (Seq.empty, M.empty)) - code1, - -- identify load_A, load_B - tmp_stms <- mapMaybe (`M.lookup` tab_inv_stm) arrs, - length tmp_stms == length arrs, - -- [inp_A, inp_B] <- arrs, - zip_AB <- zip tmp_stms arrs, - [(load_A, inp_A), (load_B, inp_B)] <- if var_dims == [0, 1] then zip_AB else reverse zip_AB, - -- code1' <- stmsFromList $ stmsToList code1 \\ stmsToList code2'', - code2' <- code2'' <> code2, - -- we get the global-thread id for the two inner dimensions, - -- as we are probably going to use it in code generation - (gtid_x, width_B) : (gtid_y, height_A) : rem_outer_dims_rev <- reverse $ unSegSpace seg_space, - rem_outer_dims <- reverse rem_outer_dims_rev, - -- sanity check that the reduce part is not missing - not $ null red_nes = do - let red_ne : _ = red_nes - red_t <- subExpType red_ne - - ---- in this binder: host code and outer seggroup (ie. the new kernel) ---- - (new_kernel, host_stms) <- runBuilder $ do - -- host code - - tk_name <- nameFromString . pretty <$> newVName "Tk" - tx_name <- nameFromString . pretty <$> newVName "Tx" - ty_name <- nameFromString . pretty <$> newVName "Ty" - rx_name <- nameFromString . pretty <$> newVName "Rx" - ry_name <- nameFromString . pretty <$> newVName "Ry" - - (ty, ry) <- getParTiles ("Ty", "Ry") (ty_name, ry_name) height_A - (tx, rx) <- getParTiles ("Tx", "Rx") (tx_name, rx_name) width_B - tk <- getSeqTile "Tk" tk_name common_dim ty tx - - tk_div_tx <- letSubExp "tk_div_tx" =<< ceilDiv tk tx - tk_div_ty <- letSubExp "tk_div_ty" =<< ceilDiv tk ty - - tx_rx <- letSubExp "TxRx" =<< toExp (pe64 tx * pe64 rx) - ty_ry <- letSubExp "TyRy" =<< toExp (pe64 ty * pe64 ry) - - a_loc_sz <- - letSubExp "a_loc_sz" - =<< toExp (pe64 ty * pe64 ry * pe64 tk) - - b_loc_sz <- - letSubExp "b_loc_sz" - =<< toExp (pe64 tk * pe64 tx * pe64 rx) - - gridDim_x <- letSubExp "gridDim_x" =<< ceilDiv width_B tx_rx - gridDim_y <- letSubExp "gridDim_y" =<< ceilDiv height_A ty_ry - let gridxy_pexp = pe64 gridDim_y * pe64 gridDim_x - let grid_pexp = - foldl (\x d -> pe64 d * x) gridxy_pexp $ - map snd rem_outer_dims_rev - grid_size <- letSubExp "grid_size" =<< toExp grid_pexp - group_size <- letSubExp "group_size" =<< toExp (pe64 ty * pe64 tx) - let segthd_lvl = SegThread (Count grid_size) (Count group_size) (SegNoVirtFull (SegSeqDims [])) - - gid_x <- newVName "gid_x" - gid_y <- newVName "gid_y" - gid_flat <- newVName "gid_flat" - - ---- in this binder: outer seggroup ---- - (ret_seggroup, stms_seggroup) <- runBuilder $ do - iii <- letExp "iii" =<< toExp (le64 gid_y * pe64 ty_ry) - jjj <- letExp "jjj" =<< toExp (le64 gid_x * pe64 tx_rx) - - -- initialize register mem with neutral elements. - cssss_list <- segMap2D "cssss" segthd_lvl ResultPrivate (ty, tx) $ \_ -> do - css_init <- scratch "css_init" (elemType red_t) [ry, rx] - css <- forLoop ry [css_init] $ \i [css_merge] -> do - css' <- forLoop rx [css_merge] $ \j [css_merge'] -> do - css'' <- update' "css" css_merge' [i, j] red_ne - resultBodyM [Var css''] - resultBodyM [Var css'] - pure [varRes css] - let [cssss] = cssss_list - - a_loc_init <- scratch "A_loc" map_t1 [a_loc_sz] - b_loc_init <- scratch "B_loc" map_t2 [b_loc_sz] - - let kkLoopBody tkind kk0 (thd_res_merge, a_loc_init', b_loc_init') epilogue = do - kk <- letExp "kk" =<< toExp (le64 kk0 * pe64 tk) - a_loc <- segScatter2D - "A_glb2loc" - a_loc_sz - a_loc_init' - segthd_lvl - [ry, tk_div_tx] - (ty, tx) - $ \[i0, k0] (thd_y, thd_x) -> do - k <- letExp "k" =<< toExp (le64 thd_x + le64 k0 * pe64 tx) - i <- letExp "i" =<< toExp (le64 thd_y + le64 i0 * pe64 ty) - - letBindNames [gtid_y] =<< toExp (le64 iii + le64 i) - a_col_idx <- letExp "A_col_idx" =<< toExp (le64 kk + le64 k) - - a_elem <- - letSubExp "A_elem" - =<< eIf - ( toExp $ - le64 gtid_y .<. pe64 height_A - .&&. if epilogue - then le64 a_col_idx .<. pe64 common_dim - else true - ) - ( do - addStm load_A - res <- index "A_elem" inp_A [a_col_idx] - resultBodyM [Var res] - ) - (eBody [eBlank $ Prim map_t1]) - a_loc_ind <- - letSubExp "a_loc_ind" - =<< eIf - (toExp $ le64 k .<. pe64 tk) - ( toExp (le64 k + le64 i * pe64 tk) - >>= letTupExp' "loc_fi" - >>= resultBodyM - ) - (eBody [pure $ BasicOp $ SubExp $ intConst Int64 (-1)]) - pure (a_elem, a_loc_ind) - - -- copy B from global to shared memory - b_loc <- segScatter2D - "B_glb2loc" - b_loc_sz - b_loc_init' - segthd_lvl - [tk_div_ty, rx] - (ty, tx) - $ \[k0, j0] (thd_y, thd_x) -> - do - k <- letExp "k" =<< toExp (le64 thd_y + le64 k0 * pe64 ty) - j <- letExp "j" =<< toExp (le64 thd_x + le64 j0 * pe64 tx) - - letBindNames [gtid_x] =<< toExp (le64 jjj + le64 j) - b_row_idx <- letExp "B_row_idx" =<< toExp (le64 kk + le64 k) - - b_elem <- - letSubExp "B_elem" - =<< eIf - ( toExp $ - le64 gtid_x .<. pe64 width_B - .&&. if epilogue - then le64 b_row_idx .<. pe64 common_dim - else true - ) - ( do - addStm load_B - res <- index "B_elem" inp_B [b_row_idx] - resultBodyM [Var res] - ) - (eBody [eBlank $ Prim map_t2]) - - b_loc_ind <- - letSubExp "b_loc_ind" - =<< eIf - (toExp $ le64 k .<. pe64 tk) - ( toExp (le64 j + le64 k * pe64 tx_rx) - >>= letTupExp' "loc_fi" - >>= resultBodyM - ) - (eBody [pure $ BasicOp $ SubExp $ intConst Int64 (-1)]) - pure (b_elem, b_loc_ind) - - -- inner loop updating this thread's accumulator (loop k in mmm_kernels). - thd_acc <- forLoop tk [thd_res_merge] $ \k [acc_merge] -> - resultBodyM =<< letTupExp' "foo" - =<< eIf - ( toExp $ - if epilogue - then - le64 kk + le64 k - .<. pe64 common_dim - else true -- if in prologue, always compute redomap. - ) - ( do - reg_mem <- segMap2D "reg_mem" segthd_lvl ResultPrivate (ty, tx) $ - \(ltid_y, ltid_x) -> do - asss_init <- scratch "asss_init" map_t1 [ry] - bsss_init <- scratch "bsss_init" map_t2 [rx] - - asss <- forLoop ry [asss_init] $ \i [asss_merge] -> do - a_loc_ind <- - letExp "a_loc_ind" - =<< toExp - ( le64 k - + (le64 ltid_y * pe64 ry + le64 i) * pe64 tk - ) - - asss <- - index "A_loc_elem" a_loc [a_loc_ind] - >>= update "asss" asss_merge [i] - resultBodyM [Var asss] - - bsss <- forLoop rx [bsss_init] $ \j [bsss_merge] -> do - b_loc_ind <- - letExp "b_loc_ind" - =<< toExp - ( le64 j - + le64 k * pe64 tx_rx - + le64 ltid_x * pe64 rx - ) - - bsss <- - index "B_loc_elem" b_loc [b_loc_ind] - >>= update "bsss" bsss_merge [j] - resultBodyM [Var bsss] - pure $ varsRes [asss, bsss] - - let [asss, bsss] = reg_mem - - -- the actual redomap. - redomap_res <- segMap2D "redomap_res" segthd_lvl ResultPrivate (ty, tx) $ - \(ltid_y, ltid_x) -> do - as <- index "as" asss [ltid_y, ltid_x] - bs <- index "bs" bsss [ltid_y, ltid_x] - css_init <- index "css_init" acc_merge [ltid_y, ltid_x] - - css <- forLoop ry [css_init] $ \i [css_merge] -> do - css <- forLoop rx [css_merge] $ \j [css_merge'] -> do - let cond = - toExp $ case tkind of - TileFull -> true - TilePartial -> - le64 iii + le64 i + pe64 ry * le64 ltid_y - .<. pe64 height_A - .&&. le64 jjj + le64 j + pe64 rx * le64 ltid_x - .<. pe64 width_B - resultBodyM =<< letTupExp' "foo" - =<< eIf - cond - ( do - a <- index "a" as [i] - b <- index "b" bs [j] - c <- index "c" css_merge' [i, j] - - map_res <- newVName "map_res" - map_lam' <- renameLambda map_lam - red_lam' <- renameLambda red_lam - - -- the inputs to map are supposed to be permutted with the - -- inverted permutation, so as to reach the original position; - -- it just so happens that the inverse of [a,b] is [b,a] - let map_inp_reg = if var_dims == [0, 1] then [a, b] else [b, a] - - addStms $ - rebindLambda map_lam' map_inp_reg [map_res] - <> rebindLambda red_lam' [c, map_res] [c] - - css <- update "css" css_merge' [i, j] c - - resultBodyM [Var css] - ) - (resultBodyM [Var css_merge']) - resultBodyM [Var css] - pure [varRes css] - - resultBodyM $ map Var redomap_res - ) - (resultBodyM [Var acc_merge]) - pure [thd_acc, a_loc, b_loc] - - -- build prologue. - full_tiles <- - letExp "full_tiles" $ - BasicOp $ BinOp (SQuot Int64 Unsafe) common_dim tk - prologue_res_list <- - forLoop' (Var full_tiles) [cssss, a_loc_init, b_loc_init] $ - \kk0 [thd_res_merge, a_loc_merge, b_loc_merge] -> do - process_full_tiles <- - kkLoopBody TileFull kk0 (thd_res_merge, a_loc_merge, b_loc_merge) False - - resultBodyM $ map Var process_full_tiles - - let prologue_res : a_loc_reuse : b_loc_reuse : _ = prologue_res_list - - -- build epilogue. - epilogue_res_list <- kkLoopBody TilePartial full_tiles (prologue_res, a_loc_reuse, b_loc_reuse) True - - let redomap_res : _ = epilogue_res_list - - -- support for non-empty code2' - -- segmap (ltid_y < ty, ltid_x < tx) { - -- for i < ry do - -- for j < rx do - -- res = if (iii+ltid_y*ry+i < height_A && jjj+ltid_x*rx+j < width_B) - -- then code2' else dummy - -- final_res[i,j] = res - epilogue_res <- - if patElemName redomap_orig_res == res_nm - then pure redomap_res -- epilogue_res_list - else do - rssss_list <- segMap2D "rssss" segthd_lvl ResultPrivate (ty, tx) $ \(ltid_y, ltid_x) -> do - rss_init <- scratch "rss_init" (elemType res_tp) [ry, rx] - css <- index "redomap_thd" redomap_res [ltid_y, ltid_x] - ii <- letExp "ii" =<< toExp (le64 iii + le64 ltid_y * pe64 ry) - jj <- letExp "jj" =<< toExp (le64 jjj + le64 ltid_x * pe64 rx) - rss <- forLoop ry [rss_init] $ \i [rss_merge] -> do - rss' <- forLoop rx [rss_merge] $ \j [rss_merge'] -> do - c <- index "redomap_elm" css [i, j] - cpy_stm <- mkLetNamesM [patElemName redomap_orig_res] $ BasicOp $ SubExp $ Var c - addStm cpy_stm - letBindNames [gtid_y] =<< toExp (le64 ii + le64 i) - letBindNames [gtid_x] =<< toExp (le64 jj + le64 j) - - res_el <- - letSubExp "res_elem" - =<< eIf - ( toExp $ - le64 gtid_y .<. pe64 height_A - .&&. le64 gtid_x .<. pe64 width_B - ) - ( do - addStms code2' - resultBodyM [Var res_nm] - ) - (eBody [eBlank res_tp]) - rss'' <- update' "rss" rss_merge' [i, j] res_el - resultBodyM [Var rss''] - resultBodyM [Var rss'] - pure [varRes rss] - let rssss : _ = rssss_list - pure rssss - - let regtile_ret_dims = - map (\(_, sz) -> (sz, se1, se1)) rem_outer_dims - ++ [(height_A, ty, ry), (width_B, tx, rx)] - - -- Add dummy dimensions to tile to reflect the outer dimensions. - epilogue_res' <- - if null rem_outer_dims - then pure epilogue_res - else do - epilogue_t <- lookupType epilogue_res - let (block_dims, rest_dims) = splitAt 2 $ arrayDims epilogue_t - ones = map (const $ intConst Int64 1) rem_outer_dims - new_shape = concat [ones, block_dims, ones, rest_dims] - letExp "res_reshaped" $ BasicOp $ Reshape (map DimNew new_shape) epilogue_res - - pure [RegTileReturns mempty regtile_ret_dims epilogue_res'] - - let level' = SegGroup (Count grid_size) (Count group_size) SegNoVirt - space' = SegSpace gid_flat (rem_outer_dims ++ [(gid_y, gridDim_y), (gid_x, gridDim_x)]) - kbody' = KernelBody () stms_seggroup ret_seggroup - pure $ Let pat aux $ Op $ SegOp $ SegMap level' space' ts kbody' - pure $ Just (host_stms, new_kernel) -mmBlkRegTiling _ = pure Nothing - -ceilDiv :: MonadBuilder m => SubExp -> SubExp -> m (Exp (Rep m)) -ceilDiv x y = pure $ BasicOp $ BinOp (SDivUp Int64 Unsafe) x y - -scratch :: MonadBuilder m => String -> PrimType -> [SubExp] -> m VName -scratch se_name t shape = letExp se_name $ BasicOp $ Scratch t shape - --- index an array with indices given in outer_indices; any inner --- dims of arr not indexed by outer_indices are sliced entirely -index :: MonadBuilder m => String -> VName -> [VName] -> m VName -index se_desc arr outer_indices = do - arr_t <- lookupType arr - let shape = arrayShape arr_t - inner_dims = shapeDims $ stripDims (length outer_indices) shape - untouched d = DimSlice (intConst Int64 0) d (intConst Int64 1) - inner_slices = map untouched inner_dims - slice = Slice $ map (DimFix . Var) outer_indices ++ inner_slices - letExp se_desc $ BasicOp $ Index arr slice - -update :: MonadBuilder m => String -> VName -> [VName] -> VName -> m VName -update se_desc arr indices new_elem = update' se_desc arr indices (Var new_elem) - -update' :: MonadBuilder m => String -> VName -> [VName] -> SubExp -> m VName -update' se_desc arr indices new_elem = - letExp se_desc $ BasicOp $ Update Unsafe arr (Slice $ map (DimFix . Var) indices) new_elem - -forLoop' :: - SubExp -> -- loop var - [VName] -> -- loop inits - ( VName -> - [VName] -> -- (loop var -> loop inits -> loop body) - Builder GPU (Body GPU) - ) -> - Builder GPU [VName] -forLoop' i_bound merge body = do - i <- newVName "i" -- could give this as arg to the function - let loop_form = ForLoop i Int64 i_bound [] - - merge_ts <- mapM lookupType merge - loop_inits <- mapM (\merge_t -> newParam "merge" $ toDecl merge_t Unique) merge_ts - - loop_body <- - runBodyBuilder . inScopeOf loop_form . localScope (scopeOfFParams loop_inits) $ - body i $ map paramName loop_inits - - letTupExp "loop" $ - DoLoop (zip loop_inits $ map Var merge) loop_form loop_body - -forLoop :: - SubExp -> - [VName] -> - (VName -> [VName] -> Builder GPU (Body GPU)) -> - Builder GPU VName -forLoop i_bound merge body = do - res_list <- forLoop' i_bound merge body - pure $ head res_list - --- given a lambda "lam", a list "new_params" of new --- parameters which should be applied to the lambda, --- and a VName "res_name" which the lambda result should --- be bound to: --- creates Stms corresponding to binding of new_params, --- lambda body, and binding of lambda result to res_name. -rebindLambda :: - Lambda GPU -> - [VName] -> - [VName] -> - Stms GPU -rebindLambda lam new_params res_names = - stmsFromList - ( zipWith - (\ident new_param -> mkLet [ident] $ BasicOp $ SubExp $ Var new_param) - idents - new_params - ) - <> bodyStms lam_body - <> stmsFromList res_cpy_stms - where - (lam_params, lam_body, lam_ret_type : _) = - (lambdaParams lam, lambdaBody lam, lambdaReturnType lam) - idents = - map - (\param -> Ident (paramName param) (paramDec param)) - lam_params - res_cpy_stms = - zipWith - ( \res_name (SubExpRes cs lam_res) -> - certify cs $ mkLet [Ident res_name lam_ret_type] $ BasicOp $ SubExp lam_res - ) - res_names - lam_ress - lam_ress = bodyResult lam_body - --- | Tries to identify the following pattern: --- code followed by some Screma followed by more code. -matchCodeStreamCode :: - Stms GPU -> - (Stms GPU, Maybe (Stm GPU), Stms GPU) -matchCodeStreamCode kstms = - let (code1, screma, code2) = - foldl - ( \acc stmt -> - case (acc, stmt) of - ((cd1, Nothing, cd2), Let _ _ (Op (OtherOp Screma {}))) -> - (cd1, Just stmt, cd2) - ((cd1, Nothing, cd2), _) -> - (cd1 ++ [stmt], Nothing, cd2) - ((cd1, Just strm, cd2), _) -> - (cd1, Just strm, cd2 ++ [stmt]) - ) - ([], Nothing, []) - (stmsToList kstms) - in (stmsFromList code1, screma, stmsFromList code2) - --- | Checks that all streamed arrays are variant to exacly one of --- the two innermost parallel dimensions, and conversely, for --- each of the two innermost parallel dimensions, there is at --- least one streamed array variant to it. The result is the --- number of the only variant parallel dimension for each array. -isInvarTo1of2InnerDims :: - Names -> - SegSpace -> - VarianceTable -> - [VName] -> - Maybe [Int] -isInvarTo1of2InnerDims branch_variant kspace variance arrs = - let inner_perm0 = map varToOnly1of2InnerDims arrs - inner_perm = catMaybes inner_perm0 - ok1 = elem 0 inner_perm && elem 1 inner_perm - ok2 = length inner_perm0 == length inner_perm - in if ok1 && ok2 then Just inner_perm else Nothing - where - varToOnly1of2InnerDims :: VName -> Maybe Int - varToOnly1of2InnerDims arr = do - (j, _) : (i, _) : _ <- Just $ reverse $ unSegSpace kspace - let variant_to = M.findWithDefault mempty arr variance - branch_invariant = - not $ nameIn j branch_variant || nameIn i branch_variant - if not branch_invariant - then Nothing -- if i or j in branch_variant; return nothing - else - if nameIn i variant_to && not (nameIn j variant_to) - then Just 0 - else - if nameIn j variant_to && not (nameIn i variant_to) - then Just 1 - else Nothing - -processIndirections :: - Names -> -- input arrays to redomap - Names -> -- variables on which the result of redomap depends on. - Maybe (Stms GPU, M.Map VName (Stm GPU)) -> - Stm GPU -> - Maybe (Stms GPU, M.Map VName (Stm GPU)) -processIndirections arrs _ acc stm@(Let patt _ (BasicOp (Index _ _))) - | Just (ss, tab) <- acc, - [p] <- patElems patt, - p_nm <- patElemName p, - nameIn p_nm arrs = - Just (ss, M.insert p_nm stm tab) -processIndirections _ res_red_var acc stm'@(Let patt _ _) - | Just (ss, tab) <- acc, - ps <- patElems patt, - all (\p -> not (nameIn (patElemName p) res_red_var)) ps = - Just (ss Seq.|> stm', tab) - | otherwise = Nothing - -se0 :: SubExp -se0 = intConst Int64 0 - -se1 :: SubExp -se1 = intConst Int64 1 - -se2 :: SubExp -se2 = intConst Int64 2 - -se4 :: SubExp -se4 = intConst Int64 4 - -se8 :: SubExp -se8 = intConst Int64 8 - -getParTiles :: (String, String) -> (Name, Name) -> SubExp -> Builder GPU (SubExp, SubExp) -getParTiles (t_str, r_str) (t_name, r_name) len_dim = - case len_dim of - Constant (IntValue (Int64Value 8)) -> - pure (se8, se1) - Constant (IntValue (Int64Value 16)) -> - pure (se8, se2) - Constant (IntValue (Int64Value 32)) -> - pure (se8, se4) - _ -> do - t <- letSubExp t_str $ Op $ SizeOp $ GetSize t_name SizeTile - r <- letSubExp r_str $ Op $ SizeOp $ GetSize r_name SizeRegTile - pure (t, r) - -getSeqTile :: String -> Name -> SubExp -> SubExp -> SubExp -> Builder GPU SubExp -getSeqTile tk_str tk_name len_dim ty tx = - case (tx, ty) of - (Constant (IntValue (Int64Value v_x)), Constant (IntValue (Int64Value v_y))) -> - letSubExp tk_str . BasicOp . SubExp . constant $ - case len_dim of - Constant (IntValue (Int64Value v_d)) -> min v_d $ min v_x v_y - _ -> min v_x v_y - _ -> - letSubExp tk_str $ Op $ SizeOp $ GetSize tk_name SizeTile - ----------------------------------------------------------------------------------------------- ---- 3D Tiling (RegTiling for the outermost dimension & Block tiling for the innermost two) --- ----------------------------------------------------------------------------------------------- - -maxRegTile :: Int64 -maxRegTile = 30 - -mkRegTileSe :: Int64 -> SubExp -mkRegTileSe = constant - -variantToDim :: VarianceTable -> VName -> VName -> Bool -variantToDim variance gid_outer nm = - gid_outer == nm || nameIn gid_outer (M.findWithDefault mempty nm variance) - --- | Checks that all streamed arrays are variant to exacly one of --- the two innermost parallel dimensions, and conversely, for --- each of the two innermost parallel dimensions, there is at --- least one streamed array variant to it. The result is the --- number of the only variant parallel dimension for each array. -isInvarTo2of3InnerDims :: - Names -> - SegSpace -> - VarianceTable -> - [VName] -> - Maybe [Int] -isInvarTo2of3InnerDims branch_variant kspace variance arrs = - let inner_perm0 = map varToOnly1of3InnerDims arrs - inner_perm = catMaybes inner_perm0 - ok1 = elem 0 inner_perm && elem 1 inner_perm && elem 2 inner_perm - ok2 = length inner_perm0 == length inner_perm - in if ok1 && ok2 then Just inner_perm else Nothing - where - varToOnly1of3InnerDims :: VName -> Maybe Int - varToOnly1of3InnerDims arr = do - (k, _) : (j, _) : (i, _) : _ <- Just $ reverse $ unSegSpace kspace - let variant_to = M.findWithDefault mempty arr variance - branch_invariant = - not $ - nameIn k branch_variant - || nameIn j branch_variant - || nameIn i branch_variant - if not branch_invariant - then Nothing -- if i or j or k in branch_variant; return nothing - else - if nameIn i variant_to && not (nameIn j variant_to || nameIn k variant_to) - then Just 0 - else - if nameIn j variant_to && not (nameIn i variant_to || nameIn k variant_to) - then Just 1 - else - if nameIn k variant_to && not (nameIn i variant_to || nameIn j variant_to) - then Just 2 - else Nothing - --- | Expects a kernel statement as argument. --- CONDITIONS for 3D tiling optimization to fire are: --- 1. a) The kernel body can be broken into --- scalar-code-1 ++ [Redomap stmt] ++ scalar-code-2. --- b) The kernels has a per-thread result, and obviously --- the result is variant to the 3rd dimension --- (counted from innermost to outermost) --- 2. For the Redomap: --- a) the streamed arrays are one dimensional --- b) each of the array arguments of Redomap are variant --- to exactly one of the three innermost-parallel dimension --- of the kernel. This condition can be relaxed by interchanging --- kernel dimensions whenever possible. --- 3. For scalar-code-1: --- a) each of the statements is a slice that produces one of the --- streamed arrays --- --- mmBlkRegTiling :: Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU)) --- mmBlkRegTiling (Let pat aux (Op (SegOp (SegMap SegThread{} seg_space ts old_kbody)))) -doRegTiling3D :: Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU)) -doRegTiling3D (Let pat aux (Op (SegOp old_kernel))) - | SegMap SegThread {} space kertp (KernelBody () kstms kres) <- old_kernel, - -- build the variance table, that records, for - -- each variable name, the variables it depends on - initial_variance <- M.map mempty $ scopeOfSegSpace space, - variance <- varianceInStms initial_variance kstms, - -- we get the global-thread id for the two inner dimensions, - -- as we are probably going to use it in code generation - (gtid_x, d_Kx) : (gtid_y, d_Ky) : (gtid_z, d_M) : rem_outer_dims_rev <- reverse $ unSegSpace space, - rem_outer_dims <- reverse rem_outer_dims_rev, - -- check that the code fits the pattern having: - -- some `code1`, followed by one Screma SOAC, followed by some `code2` - (code1, Just screma_stmt, code2) <- matchCodeStreamCode kstms, - Let pat_redomap _ (Op _) <- screma_stmt, - -- checks that the Screma SOAC is actually a redomap and normalize it - Just (common_dim, inp_soac_arrs, (_, red_lam, red_nes, map_lam)) <- isTileableRedomap screma_stmt, - not (null red_nes), - -- assuming we have a budget of maxRegTile registers, we distribute - -- that budget across the result of redomap and the kernel result - num_res <- max (length red_nes) (length kres), - reg_tile <- maxRegTile `quot` fromIntegral num_res, - reg_tile_se <- mkRegTileSe reg_tile, - -- check that the element-type of the map and reduce are scalars: - all (primType . paramDec) $ lambdaParams map_lam, - red_res_tps <- map paramDec $ take (length red_nes) $ lambdaParams red_lam, - all primType red_res_tps, - -- checks that the input arrays to redomap are variant to - -- exactly one of the two innermost dimensions of the kernel - Just _ <- isInvarTo2of3InnerDims mempty space variance inp_soac_arrs, - -- get the free variables on which the result of redomap depends on - redomap_orig_res <- patElems pat_redomap, - res_red_var <- -- variance of the reduce result - mconcat $ mapMaybe ((`M.lookup` variance) . patElemName) redomap_orig_res, - mempty /= res_red_var, - -- we furthermore check that code1 is only formed by - -- 1. statements that slice some globally-declared arrays - -- to produce the input for the redomap, and - -- 2. potentially some statements on which the redomap - -- is independent; these are recorded in `code2''` - Just (code2'', arr_tab0) <- - foldl - (processIndirections (namesFromList inp_soac_arrs) res_red_var) - (Just (Seq.empty, M.empty)) - code1, - -- check that code1 contains exacly one slice for each of the input array to redomap - tmp_stms <- mapMaybe (`M.lookup` arr_tab0) inp_soac_arrs, - length tmp_stms == length inp_soac_arrs, - -- code1' <- stmsFromList $ stmsToList code1 \\ stmsToList code2'', - code2' <- code2'' <> code2, - -- we assume the kernel results are variant to the thrid-outer parallel dimension - -- (for sanity sake, they should be) - ker_res_nms <- mapMaybe getResNm kres, - length ker_res_nms == length kres, - all primType kertp, - all (variantToDim variance gtid_z) ker_res_nms = do - -- HERE STARTS THE IMPLEMENTATION: - (new_kernel, host_stms) <- runBuilder $ do - -- host code - -- process the z-variant arrays that need transposition; - -- these "manifest" statements should come before the kernel - (tab_inn, tab_out) <- - foldM - (insertTranspose variance (gtid_z, d_M)) - (M.empty, M.empty) - $ M.toList arr_tab0 - - tx_name <- nameFromString . pretty <$> newVName "Tx" - ty_name <- nameFromString . pretty <$> newVName "Ty" - - tx0 <- letSubExp "Tx" $ Op $ SizeOp $ GetSize tx_name SizeTile - ty0 <- letSubExp "Ty" $ Op $ SizeOp $ GetSize ty_name SizeTile - ty <- limitTile "Ty" ty0 d_Ky - tx <- limitTile "Tx" tx0 d_Kx - let rz = reg_tile_se - - gridDim_x <- letSubExp "gridDim_x" =<< ceilDiv d_Kx tx - gridDim_y <- letSubExp "gridDim_y" =<< ceilDiv d_Ky ty - gridDim_z <- letSubExp "gridDim_z" =<< ceilDiv d_M rz - let gridxyz_pexp = pe64 gridDim_z * pe64 gridDim_y * pe64 gridDim_x - let grid_pexp = product $ gridxyz_pexp : map (pe64 . snd) rem_outer_dims_rev - grid_size <- letSubExp "grid_size_tile3d" =<< toExp grid_pexp - group_size <- letSubExp "group_size_tile3d" =<< toExp (pe64 ty * pe64 tx) - let segthd_lvl = SegThread (Count grid_size) (Count group_size) (SegNoVirtFull (SegSeqDims [])) - - count_shmem <- letSubExp "count_shmem" =<< ceilDiv rz group_size - - gid_x <- newVName "gid_x" - gid_y <- newVName "gid_y" - gid_z <- newVName "gid_z" - gid_flat <- newVName "gid_flat" - - ---- in this binder: outer seggroup ---- - (ret_seggroup, stms_seggroup) <- runBuilder $ do - ii <- letExp "ii" =<< toExp (le64 gid_z * pe64 rz) - jj1 <- letExp "jj1" =<< toExp (le64 gid_y * pe64 ty) - jj2 <- letExp "jj2" =<< toExp (le64 gid_x * pe64 tx) - - -- initialize the register arrays corresponding to the result of redomap; - reg_arr_nms <- segMap2D "res" segthd_lvl ResultPrivate (ty, tx) $ \_ -> - forM (zip red_nes red_res_tps) $ \(red_ne, red_t) -> do - css_init <- scratch "res_init" (elemType red_t) [rz] - css <- forLoop rz [css_init] $ \i [css_merge] -> do - css' <- update' "css" css_merge [i] red_ne - resultBodyM [Var css'] - pure $ varRes css - - -- scratch the shared-memory arrays corresponding to the arrays that are - -- input to the redomap and are invariant to the outermost parallel dimension. - loc_arr_nms <- forM (M.toList tab_out) $ \(nm, (ptp, _)) -> - scratch (baseString nm ++ "_loc") ptp [rz] - - prologue_res_list <- - forLoop' common_dim (reg_arr_nms ++ loc_arr_nms) $ - \q var_nms -> do - let reg_arr_merge_nms = take (length red_nes) var_nms - let loc_arr_merge_nms = drop (length red_nes) var_nms - - -- collective copy from global to shared memory - loc_arr_nms' <- - forLoop' count_shmem loc_arr_merge_nms $ \tt loc_arr_merge2_nms -> do - loc_arr_merge2_nms' <- - forM (zip loc_arr_merge2_nms (M.toList tab_out)) $ \(loc_Y_nm, (glb_Y_nm, (ptp_Y, load_Y))) -> do - ltid_flat <- newVName "ltid_flat" - ltid <- newVName "ltid" - let segspace = SegSpace ltid_flat [(ltid, group_size)] - ((res_v, res_i), stms) <- runBuilder $ do - offs <- letExp "offs" =<< toExp (pe64 group_size * le64 tt) - loc_ind <- letExp "loc_ind" =<< toExp (le64 ltid + le64 offs) - letBindNames [gtid_z] =<< toExp (le64 ii + le64 loc_ind) - let glb_ind = gtid_z - y_elm <- - letSubExp "y_elem" - =<< eIf - (toExp $ le64 glb_ind .<. pe64 d_M) - ( do - addStm load_Y - res <- index "Y_elem" glb_Y_nm [q] - resultBodyM [Var res] - ) - (eBody [eBlank $ Prim ptp_Y]) - y_ind <- - letSubExp "y_loc_ind" - =<< eIf - (toExp $ le64 loc_ind .<. pe64 rz) - (toExp loc_ind >>= letTupExp' "loc_fi" >>= resultBodyM) - (eBody [pure $ BasicOp $ SubExp $ intConst Int64 (-1)]) - -- y_tp <- subExpType y_elm - pure (y_elm, y_ind) - - let ret = WriteReturns mempty (Shape [rz]) loc_Y_nm [(Slice [DimFix res_i], res_v)] - let body = KernelBody () stms [ret] - - res_nms <- - letTupExp "Y_glb2loc" <=< renameExp $ - Op $ SegOp $ SegMap segthd_lvl segspace [Prim ptp_Y] body - let res_nm : _ = res_nms - pure res_nm - resultBodyM $ map Var loc_arr_merge2_nms' - - redomap_res <- - segMap2D "redomap_res" segthd_lvl ResultPrivate (ty, tx) $ - \(ltid_y, ltid_x) -> do - letBindNames [gtid_y] =<< toExp (le64 jj1 + le64 ltid_y) - letBindNames [gtid_x] =<< toExp (le64 jj2 + le64 ltid_x) - reg_arr_merge_nms_slc <- forM reg_arr_merge_nms $ \reg_arr_nm -> - index "res_reg_slc" reg_arr_nm [ltid_y, ltid_x] - fmap subExpsRes . letTupExp' "redomap_guarded" - =<< eIf - (toExp $ le64 gtid_y .<. pe64 d_Ky .&&. le64 gtid_x .<. pe64 d_Kx) - ( do - inp_scals_invar_outer <- - forM (M.toList tab_inn) $ \(inp_arr_nm, load_stm) -> do - addStm load_stm - index (baseString inp_arr_nm) inp_arr_nm [q] - -- build the loop of count R whose body is semantically the redomap code - reg_arr_merge_nms' <- - forLoop' rz reg_arr_merge_nms_slc $ \i reg_arr_mm_nms -> do - letBindNames [gtid_z] =<< toExp (le64 ii + le64 i) - resultBodyM =<< letTupExp' "redomap_lam" - =<< eIf - (toExp $ le64 gtid_z .<. pe64 d_M) - ( do - -- read from shared memory - ys <- forM loc_arr_nms' $ \loc_arr_nm -> - index "inp_reg_var2z" loc_arr_nm [i] - cs <- forM reg_arr_mm_nms $ \reg_arr_nm -> - index "res_reg_var2z" reg_arr_nm [i] - -- here we need to put in order the scalar inputs to map: - let tab_scals = - M.fromList $ - zip (map fst $ M.toList tab_out) ys - ++ zip (map fst $ M.toList tab_inn) inp_scals_invar_outer - map_inp_scals <- forM inp_soac_arrs $ \arr_nm -> - case M.lookup arr_nm tab_scals of - Nothing -> error "Impossible case reached in tiling3D\n" - Just nm -> pure nm - map_res_scals <- forM (lambdaReturnType map_lam) $ \_ -> newVName "map_res" - map_lam' <- renameLambda map_lam - red_lam' <- renameLambda red_lam - addStms $ - rebindLambda map_lam' map_inp_scals map_res_scals - <> rebindLambda red_lam' (cs ++ map_res_scals) cs - css <- forM (zip reg_arr_mm_nms cs) $ \(reg_arr_nm, c) -> - update (baseString reg_arr_nm) reg_arr_nm [i] c - resultBodyM $ map Var css - ) - (resultBodyM $ map Var reg_arr_mm_nms) - resultBodyM $ map Var reg_arr_merge_nms' - ) - (resultBodyM $ map Var reg_arr_merge_nms_slc) - resultBodyM $ map Var $ redomap_res ++ loc_arr_nms' - - -- support for non-empty code2' - -- segmap (ltid_y < ty, ltid_x < tx) { - -- for i < rz do - -- res = if (ii+i < d_M && jj1+ltid_y < d_Ky && jj2 + ltid_x < d_Kx) - -- then code2' else dummy - -- final_res[i] = res - let redomap_res = take (length red_nes) prologue_res_list - epilogue_res <- - if length redomap_orig_res == length ker_res_nms - && ker_res_nms == map patElemName redomap_orig_res - then -- all (\ (a,b) -> patElemName a == b ) $ zip redomap_orig_res ker_res_nms - segMap3D "rssss" segthd_lvl ResultPrivate (se1, ty, tx) $ \(_ltid_z, ltid_y, ltid_x) -> +import Data.List.NonEmpty (NonEmpty (..)) +import qualified Data.Map.Strict as M +import Data.Maybe +import qualified Data.Sequence as Seq +import Futhark.IR.GPU +import qualified Futhark.IR.Mem.IxFun as IxFun +import Futhark.MonadFreshNames +import Futhark.Optimise.TileLoops.Shared +import Futhark.Tools +import Futhark.Transform.Rename +import Futhark.Transform.Substitute + +se0 :: SubExp +se0 = intConst Int64 0 + +se1 :: SubExp +se1 = intConst Int64 1 + +se2 :: SubExp +se2 = intConst Int64 2 + +se4 :: SubExp +se4 = intConst Int64 4 + +se8 :: SubExp +se8 = intConst Int64 8 + +scratch :: MonadBuilder m => String -> PrimType -> [SubExp] -> m VName +scratch se_name t shape = letExp se_name $ BasicOp $ Scratch t shape + +-- | Main helper function for Register-and-Block Tiling +kkLoopBody :: + Env -> + ( (SubExp, SubExp, SubExp, SubExp, SubExp, SubExp, SubExp, SubExp), + SegLevel, + [Int], + (VName, SubExp, VName, SubExp, SubExp), + (SubExp, SubExp), + (VName, VName), + (Stm GPU, VName, PrimType, Stm GPU, VName, PrimType), + (Lambda GPU, Lambda GPU) + ) -> + VName -> + (VName, VName, VName) -> + Bool -> + Builder GPU [VName] +kkLoopBody + env + ( (rx, ry, tx, ty, tk, tk_div_tx, _tk_div_ty, tx_rx), + segthd_lvl, + var_dims, + (gtid_x, width_B, gtid_y, height_A, common_dim), + (a_loc_sz, b_loc_sz), + (iii, jjj), + (load_A, inp_A, pt_A, load_B, inp_B, pt_B), + (map_lam, red_lam) + ) + kk0 + (thd_res_merge, a_loc_init', b_loc_init') + epilogue = do + let (map_t1, map_t2) = (pt_A, pt_B) + kk <- letExp "kk" =<< toExp (le64 kk0 * pe64 tk) + -- copy A to local memory + (a_loc, aCopyLoc2Reg) <- + copyGlb2ShMem kk (gtid_y, iii, map_t1, height_A, inp_A, load_A, a_loc_sz, a_loc_init') + + -- copy B from global to shared memory + (b_loc, bCopyLoc2Reg) <- + copyGlb2ShMem kk (gtid_x, jjj, map_t2, width_B, inp_B, load_B, b_loc_sz, b_loc_init') + + -- inner loop updating this thread's accumulator (loop k in mmm_kernels). + thd_acc <- forLoop tk [thd_res_merge] $ \k [acc_merge] -> + resultBodyM =<< letTupExp' "foo" + =<< eIf + ( toExp $ + if epilogue + then le64 kk + le64 k .<. pe64 common_dim + else true -- if in prologue, always compute redomap. + ) + ( do + reg_mem <- segMap2D "reg_mem" segthd_lvl ResultPrivate (ty, tx) $ + \(ltid_y, ltid_x) -> do + -- copy A from local memory to registers + asss <- aCopyLoc2Reg k ltid_y + -- copy B from local memory to registers + bsss <- bCopyLoc2Reg k ltid_x + pure $ varsRes [asss, bsss] + let [asss, bsss] = reg_mem + mkRedomapOneTileBody acc_merge asss bsss True + ) + (resultBodyM [Var acc_merge]) + pure [thd_acc, a_loc, b_loc] + where + mk_ik is_coal (thd_y, thd_x) (i0, k0) + | is_coal = do + -- not-transposed case (i.e., already coalesced) + let (t_par, t_seq) = (tx, tk) + k <- letExp "k" =<< toExp (le64 thd_x + le64 k0 * pe64 t_par) + i <- letExp "i" =<< toExp (le64 thd_y + le64 i0 * pe64 t_par) + -- we have padded to minimize bank conflicts, + -- hence the length of inner dim is (t_seq + 1) + let e = le64 k + le64 i * (pe64 t_seq + pe64 se1) + pure (i, k, e) + mk_ik _ (thd_y, thd_x) (i0, k0) = do + -- matrix is transposed case (i.e., uncoalesced): + let (t_par, tr_par) = (tx, tx_rx) + k <- letExp "k" =<< toExp (le64 thd_y + le64 k0 * pe64 t_par) + i <- letExp "i" =<< toExp (le64 thd_x + le64 i0 * pe64 t_par) + -- we have padded to minimize bank conflicts, + -- hence the length of inner dim is (tr_par + 1) + let e = le64 i + le64 k * (pe64 tr_par + pe64 se1) + pure (i, k, e) + isInnerCoal :: Env -> VName -> Stm GPU -> Bool + isInnerCoal (_, ixfn_env) slc_X (Let pat _ (BasicOp (Index x _))) + | [slc_X'] <- patNames pat, + slc_X == slc_X', + Nothing <- M.lookup x ixfn_env = + True -- if not in the table, we assume not-transposed! + isInnerCoal (_, ixfn_env) slc_X (Let pat _ (BasicOp (Index x _))) + | [slc_X'] <- patNames pat, + slc_X == slc_X', + Just ixf_fn <- M.lookup x ixfn_env, + (IxFun.IxFun (lmad :| []) _ _) <- ixf_fn = + let lmad_dims = IxFun.lmadDims lmad + q = length lmad_dims + last_perm = IxFun.ldPerm $ last lmad_dims + stride = IxFun.ldStride $ last lmad_dims + res = last_perm == q - 1 && (stride == pe64 (intConst Int64 1)) + in res + isInnerCoal _ _ _ = error "TileLoops/Shared.hs: not an error, but I would like to know why!" + -- + mkRedomapOneTileBody acc_merge asss bsss fits_ij = do + -- the actual redomap. + redomap_res <- segMap2D "redomap_res" segthd_lvl ResultPrivate (ty, tx) $ + \(ltid_y, ltid_x) -> do + as <- index "as" asss [ltid_y, ltid_x] + bs <- index "bs" bsss [ltid_y, ltid_x] + css_init <- index "css_init" acc_merge [ltid_y, ltid_x] + + css <- forLoop ry [css_init] $ \i [css_merge] -> do + css <- forLoop rx [css_merge] $ \j [css_merge'] -> + resultBodyM =<< letTupExp' "foo" + =<< eIf + ( toExp $ + if fits_ij + then true + else -- this condition is never needed because + -- if i and j are out of range than css[i,j] + -- is garbage anyways and should not be written. + -- so fits_ij should be always true!!! + + le64 iii + le64 i + pe64 ry * le64 ltid_y + .<. pe64 height_A + .&&. le64 jjj + le64 j + pe64 rx * le64 ltid_x + .<. pe64 width_B + ) + ( do + a <- index "a" as [i] + b <- index "b" bs [j] + c <- index "c" css_merge' [i, j] + + map_lam' <- renameLambda map_lam + red_lam' <- renameLambda red_lam + + -- the inputs to map are supposed to be permutted with the + -- inverted permutation, so as to reach the original position; + -- it just so happens that the inverse of [a,b] is [b,a] + let map_inp_reg = if var_dims == [0, 1] then [a, b] else [b, a] + + map_res <- eLambda map_lam' (map (eSubExp . Var) map_inp_reg) + ~[red_res] <- eLambda red_lam' (map eSubExp $ Var c : map resSubExp map_res) + css <- update "css" css_merge' [i, j] (resSubExp red_res) + + resultBodyM [Var css] + ) + (resultBodyM [Var css_merge']) + resultBodyM [Var css] + pure [varRes css] + resultBodyM $ map Var redomap_res + -- + copyGlb2ShMem :: + VName -> + (VName, VName, PrimType, SubExp, VName, Stm GPU, SubExp, VName) -> + Builder GPU (VName, VName -> VName -> Builder GPU VName) + copyGlb2ShMem kk (gtid, ii, ptp_X_el, parlen_X, inp_X, load_X, loc_sz_X, x_loc_init') = do + let (t_par, r_par, tseq_div_tpar) = (tx, rx, tk_div_tx) + is_inner_coal = isInnerCoal env inp_X load_X + str_A = baseString inp_X + x_loc <- + segScatter2D (str_A ++ "_glb2loc") loc_sz_X x_loc_init' segthd_lvl [r_par, tseq_div_tpar] (t_par, t_par) $ + scatterFun is_inner_coal + + pure (x_loc, copyLoc2Reg is_inner_coal str_A x_loc) + where + copyLoc2Reg :: + Bool -> + String -> + VName -> + VName -> + VName -> + Builder GPU VName + copyLoc2Reg is_inner_coal str_A x_loc k ltid_yx = do + let (r_par, t_seq, tr_par) = (rx, tk, tx_rx) + xsss_init <- scratch (str_A ++ "_init_regs") ptp_X_el [r_par] + forLoop r_par [xsss_init] $ \ij [xsss_merge] -> do + x_loc_ind <- + letExp (str_A ++ "_loc_ind") + =<< toExp + ( if is_inner_coal + then le64 k + (le64 ltid_yx * pe64 r_par + le64 ij) * (pe64 t_seq + pe64 se1) + else le64 ij + le64 ltid_yx * pe64 r_par + le64 k * (pe64 tr_par + pe64 se1) + ) + xsss <- + update (str_A ++ "_regs") xsss_merge [ij] . Var + =<< index (str_A ++ "_loc_elem") x_loc [x_loc_ind] + resultBodyM [Var xsss] + -- + scatterFun :: + Bool -> + [VName] -> + (VName, VName) -> + Builder GPU (SubExp, SubExp) + scatterFun is_inner_coal [i0, k0] (thd_y, thd_x) = do + let str_A = baseString inp_X + t_seq = tk + (i, k, epx_loc_fi) <- mk_ik is_inner_coal (thd_y, thd_x) (i0, k0) + letBindNames [gtid] =<< toExp (le64 ii + le64 i) + a_seqdim_idx <- letExp (str_A ++ "_seqdim_idx") =<< toExp (le64 kk + le64 k) + + a_elem <- + letSubExp (str_A ++ "_elem") + =<< eIf + ( toExp $ + le64 gtid .<. pe64 parlen_X + .&&. if epilogue + then le64 a_seqdim_idx .<. pe64 common_dim + else true + ) + ( do + addStm load_X + res <- index "A_elem" inp_X [a_seqdim_idx] + resultBodyM [Var res] + ) + (eBody [eBlank $ Prim ptp_X_el]) + + a_loc_ind <- + letSubExp (str_A ++ "_loc_ind") + =<< eIf + (toExp $ le64 k .<. pe64 t_seq) + (eBody [toExp epx_loc_fi]) + (eBody [eSubExp $ intConst Int64 (-1)]) + pure (a_elem, a_loc_ind) + scatterFun _ _ _ = do + error "Function scatterFun in Shared.hs: 2nd arg should be an array with 2 elements!" + +-- ToDo: we need tx == ty (named t_par), and rx == ry (named r_par) +-- in order to handle all the cases without transpositions. +-- additionally, of course, we need that tk is a multiple of t_par. +mmBlkRegTiling :: Env -> Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU)) +mmBlkRegTiling env stm = do + res <- mmBlkRegTilingAcc env stm + case res of + Nothing -> mmBlkRegTilingNrm env stm + _ -> pure res + +mmBlkRegTilingAcc :: Env -> Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU)) +mmBlkRegTilingAcc env (Let pat aux (Op (SegOp (SegMap SegThread {} seg_space ts old_kbody)))) + | KernelBody () kstms [Returns ResultMaySimplify cs (Var res_nm)] <- old_kbody, + cs == mempty, + -- check kernel has one result of primitive type + [res_tp] <- ts, + isAcc res_tp, + -- we get the global-thread id for the two inner dimensions, + -- as we are probably going to use it in code generation + (gtid_x, width_B) : (gtid_y, height_A) : rem_outer_dims_rev <- + reverse $ unSegSpace seg_space, + rem_outer_dims <- reverse rem_outer_dims_rev, + Just + ( code2', + (load_A, inp_A, map_t1, load_B, inp_B, map_t2), + common_dim, + var_dims, + (map_lam, red_lam, red_ne, redomap_orig_res, red_t) + ) <- + matchesBlkRegTile seg_space kstms, + checkAccumulatesRedomapRes res_nm code2' redomap_orig_res = do + -- Here we start the implementation -- + ---- in this binder: host code and outer seggroup (ie. the new kernel) ---- + (new_kernel, host_stms) <- runBuilder $ do + -- host code + (rx, ry, tx, ty, tk, tk_div_tx, tk_div_ty, tx_rx, ty_ry, a_loc_sz, b_loc_sz) <- + mkTileMemSizes height_A width_B common_dim + + rk <- letSubExp "rk" $ BasicOp $ SubExp $ intConst Int64 8 -- 16 and 8 seem good values + tk_rk <- letSubExp "tk_rk" =<< toExp (pe64 tk * pe64 rk) + + gridDim_t <- letSubExp "gridDim_t" =<< ceilDiv common_dim tk_rk + gridDim_y <- letSubExp "gridDim_y" =<< ceilDiv height_A ty_ry + gridDim_x <- letSubExp "gridDim_x" =<< ceilDiv width_B tx_rx + + let gridxyt_pexp = pe64 gridDim_y * pe64 gridDim_x * pe64 gridDim_t + grid_pexp = + foldl (\x d -> pe64 d * x) gridxyt_pexp $ + map snd rem_outer_dims_rev + + (grid_size, group_size, segthd_lvl) <- mkNewSegthdLvl tx ty grid_pexp + (gid_x, gid_y, gid_flat) <- mkGidsXYF + gid_t <- newVName "gid_t" + + ---- in this binder: outer seggroup ---- + (ret_seggroup, stms_seggroup) <- runBuilder $ do + iii <- letExp "iii" =<< toExp (le64 gid_y * pe64 ty_ry) + jjj <- letExp "jjj" =<< toExp (le64 gid_x * pe64 tx_rx) + ttt <- letExp "ttt" =<< toExp (le64 gid_t * pe64 tk_rk) + + -- initialize register mem with neutral elements and create shmem + (cssss, a_loc_init, b_loc_init) <- + initRegShmem + (rx, tx, ry, ty, a_loc_sz, b_loc_sz) + (map_t1, map_t2, red_t) + segthd_lvl + red_ne + + -- build prologue. + elems_on_t <- letSubExp "elems_on_t" =<< toExp (pe64 common_dim - le64 ttt) + tiles_on_t <- letSubExp "tiles_on_t" $ BasicOp $ BinOp (SQuot Int64 Unsafe) elems_on_t tk + full_tiles <- letExp "full_tiles" $ BasicOp $ BinOp (SMin Int64) rk tiles_on_t + + let ct_arg = + ( (rx, ry, tx, ty, tk, tk_div_tx, tk_div_ty, tx_rx), + segthd_lvl, + var_dims, + (gtid_x, width_B, gtid_y, height_A, common_dim), + (a_loc_sz, b_loc_sz), + (iii, jjj), + (load_A, inp_A, map_t1, load_B, inp_B, map_t2), + (map_lam, red_lam) + ) + + prologue_res_list <- + forLoop' (Var full_tiles) [cssss, a_loc_init, b_loc_init] $ + \kk0 [thd_res_merge, a_loc_merge, b_loc_merge] -> do + off_t <- letExp "off_t" =<< toExp (pe64 rk * le64 gid_t + le64 kk0) + process_full_tiles <- + kkLoopBody env ct_arg off_t (thd_res_merge, a_loc_merge, b_loc_merge) False + + resultBodyM $ map Var process_full_tiles + + let prologue_res : a_loc_reuse : b_loc_reuse : _ = prologue_res_list + + redomap_res_lst <- + letTupExp "redomap_res_if" + =<< eIf + ( toExp $ + le64 full_tiles .==. pe64 rk + .||. pe64 common_dim .==. (pe64 tk * le64 full_tiles + le64 ttt) + ) + (resultBodyM $ map Var prologue_res_list) + ( do + off_t <- letExp "off_t" =<< toExp (pe64 rk * le64 gid_t + le64 full_tiles) + process_sprs_tile <- + kkLoopBody env ct_arg off_t (prologue_res, a_loc_reuse, b_loc_reuse) True + + resultBodyM $ map Var process_sprs_tile + ) + let redomap_res : _ = redomap_res_lst + + -- support for non-empty code2' + -- segmap (ltid_y < ty, ltid_x < tx) { + -- for i < ry do + -- for j < rx do + -- res = if (iii+ltid_y*ry+i < height_A && jjj+ltid_x*rx+j < width_B) + -- then code2' else dummy + -- final_res[i,j] = res + mkEpilogueAccRes + segthd_lvl + (redomap_orig_res, redomap_res) + (res_nm, res_tp) + (ty, tx, ry, rx) + (iii, jjj) + (gtid_y, gtid_x) + (height_A, width_B, rem_outer_dims) + code2' + + let level' = SegGroup (Count grid_size) (Count group_size) SegNoVirt + space' = SegSpace gid_flat (rem_outer_dims ++ [(gid_t, gridDim_t), (gid_y, gridDim_y), (gid_x, gridDim_x)]) + kbody' = KernelBody () stms_seggroup ret_seggroup + pure $ Let pat aux $ Op $ SegOp $ SegMap level' space' ts kbody' + pure $ Just (host_stms, new_kernel) + where + sameAccType acc_sglton (Acc sglton _ _ _) = + acc_sglton == sglton + sameAccType _ _ = False + getAccumFV (Acc singleton _shp [_eltp] _) = do + let fvs = namesToList $ freeIn old_kbody -- code + tps <- localScope (scopeOfSegSpace seg_space) $ do + mapM lookupType fvs + let (acc_0s, _) = unzip $ filter (sameAccType singleton . snd) $ zip fvs tps + case acc_0s of + [acc_0] -> pure acc_0 + _ -> error "Impossible case reached when treating accumulators!" + getAccumFV tp = error ("Should be an accumulator type at this point, given: " ++ pretty tp) + -- + -- checks that the redomap result is used directly as the accumulated value, + -- in which case it is safe to parallelize the innermost dimension (of tile tk) + checkAccumulatesRedomapRes res_nm acc_code redomap_orig_res = do + foldl getAccumStm False $ reverse $ stmsToList acc_code + where + getAccumStm True _ = True + getAccumStm False (Let (Pat [pat_el]) _aux (BasicOp (UpdateAcc _acc_nm _ind vals))) + | [v] <- vals, + patElemName pat_el == res_nm = + v == Var redomap_orig_res + getAccumStm False _ = False + -- + -- epilogue for accumulator result type + mkEpilogueAccRes + segthd_lvl + (redomap_orig_res, redomap_res) + (res_nm, res_tp) + (ty, tx, ry, rx) + (iii, jjj) + (gtid_y, gtid_x) + (height_A, width_B, _rem_outer_dims) + code2' = do + rss_init <- getAccumFV res_tp + rssss_list <- segMap2D "rssss" segthd_lvl ResultMaySimplify (ty, tx) $ \(ltid_y, ltid_x) -> do + (css, ii, jj) <- getThdRedomapRes (rx, ry) (ltid_x, ltid_y) (iii, jjj, redomap_res) + rss <- forLoop ry [rss_init] $ \i [rss_merge] -> do + rss' <- forLoop rx [rss_merge] $ \j [rss_merge'] -> do + prereqAddCode2 (gtid_x, gtid_y) (ii, i, jj, j) (css, redomap_orig_res) + let code2_subs = substituteNames (M.singleton rss_init rss_merge') code2' + + res_el <- + letSubExp "res_elem" + =<< eIf + ( toExp $ + le64 gtid_y .<. pe64 height_A + .&&. le64 gtid_x .<. pe64 width_B + ) + ( do + addStms code2_subs + resultBodyM [Var res_nm] + ) + (resultBodyM [Var rss_merge']) + resultBodyM [res_el] + resultBodyM [Var rss'] + pure [varRes rss] + let epilogue_res_acc : _ = rssss_list + pure [Returns ResultMaySimplify (Certs []) $ Var epilogue_res_acc] +mmBlkRegTilingAcc _ _ = pure Nothing + +-------------------------- +-------------------------- + +mmBlkRegTilingNrm :: Env -> Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU)) +mmBlkRegTilingNrm env (Let pat aux (Op (SegOp (SegMap SegThread {} seg_space ts old_kbody)))) + | KernelBody () kstms [Returns ResultMaySimplify cs (Var res_nm)] <- old_kbody, + cs == mempty, + -- check kernel has one result of primitive type + [res_tp] <- ts, + primType res_tp, + -- we get the global-thread id for the two inner dimensions, + -- as we are probably going to use it in code generation + (gtid_x, width_B) : (gtid_y, height_A) : rem_outer_dims_rev <- + reverse $ unSegSpace seg_space, + rem_outer_dims <- reverse rem_outer_dims_rev, + Just + ( code2', + (load_A, inp_A, map_t1, load_B, inp_B, map_t2), + common_dim, + var_dims, + (map_lam, red_lam, red_ne, redomap_orig_res, red_t) + ) <- + matchesBlkRegTile seg_space kstms = do + -- Here we start the implementation + ---- in this binder: host code and outer seggroup (ie. the new kernel) ---- + (new_kernel, host_stms) <- runBuilder $ do + -- host code + (rx, ry, tx, ty, tk, tk_div_tx, tk_div_ty, tx_rx, ty_ry, a_loc_sz, b_loc_sz) <- + mkTileMemSizes height_A width_B common_dim + + gridDim_x <- letSubExp "gridDim_x" =<< ceilDiv width_B tx_rx + gridDim_y <- letSubExp "gridDim_y" =<< ceilDiv height_A ty_ry + let gridxy_pexp = pe64 gridDim_y * pe64 gridDim_x + let grid_pexp = + foldl (\x d -> pe64 d * x) gridxy_pexp $ + map snd rem_outer_dims_rev + (grid_size, group_size, segthd_lvl) <- mkNewSegthdLvl tx ty grid_pexp + + (gid_x, gid_y, gid_flat) <- mkGidsXYF + + ---- in this binder: outer seggroup ---- + (ret_seggroup, stms_seggroup) <- runBuilder $ do + iii <- letExp "iii" =<< toExp (le64 gid_y * pe64 ty_ry) + jjj <- letExp "jjj" =<< toExp (le64 gid_x * pe64 tx_rx) + + -- initialize register mem with neutral elements and create shmem + (cssss, a_loc_init, b_loc_init) <- + initRegShmem + (rx, tx, ry, ty, a_loc_sz, b_loc_sz) + (map_t1, map_t2, red_t) + segthd_lvl + red_ne + + -- build prologue. + full_tiles <- + letExp "full_tiles" $ + BasicOp $ BinOp (SQuot Int64 Unsafe) common_dim tk + + let ct_arg = + ( (rx, ry, tx, ty, tk, tk_div_tx, tk_div_ty, tx_rx), + segthd_lvl, + var_dims, + (gtid_x, width_B, gtid_y, height_A, common_dim), + (a_loc_sz, b_loc_sz), + (iii, jjj), + (load_A, inp_A, map_t1, load_B, inp_B, map_t2), + (map_lam, red_lam) + ) + + prologue_res_list <- + forLoop' (Var full_tiles) [cssss, a_loc_init, b_loc_init] $ + \kk0 [thd_res_merge, a_loc_merge, b_loc_merge] -> do + process_full_tiles <- + kkLoopBody env ct_arg kk0 (thd_res_merge, a_loc_merge, b_loc_merge) False + + resultBodyM $ map Var process_full_tiles + + let prologue_res : a_loc_reuse : b_loc_reuse : _ = prologue_res_list + + -- build epilogue. + epilogue_res_list <- kkLoopBody env ct_arg full_tiles (prologue_res, a_loc_reuse, b_loc_reuse) True + + let redomap_res : _ = epilogue_res_list + + -- support for non-empty code2' + -- segmap (ltid_y < ty, ltid_x < tx) { + -- for i < ry do + -- for j < rx do + -- res = if (iii+ltid_y*ry+i < height_A && jjj+ltid_x*rx+j < width_B) + -- then code2' else dummy + -- final_res[i,j] = res + mkEpiloguePrimRes + segthd_lvl + (redomap_orig_res, redomap_res) + (res_nm, res_tp) + (ty, tx, ry, rx) + (iii, jjj) + (gtid_y, gtid_x) + (height_A, width_B, rem_outer_dims) + code2' + + let level' = SegGroup (Count grid_size) (Count group_size) SegNoVirt + space' = SegSpace gid_flat (rem_outer_dims ++ [(gid_y, gridDim_y), (gid_x, gridDim_x)]) + kbody' = KernelBody () stms_seggroup ret_seggroup + pure $ Let pat aux $ Op $ SegOp $ SegMap level' space' ts kbody' + pure $ Just (host_stms, new_kernel) + where + mkEpiloguePrimRes + segthd_lvl + (redomap_orig_res, redomap_res) + (res_nm, res_tp) + (ty, tx, ry, rx) + (iii, jjj) + (gtid_y, gtid_x) + (height_A, width_B, rem_outer_dims) + code2' = do + epilogue_res <- + if redomap_orig_res == res_nm + then pure redomap_res -- epilogue_res_list + else do + rssss_list <- segMap2D "rssss" segthd_lvl ResultPrivate (ty, tx) $ \(ltid_y, ltid_x) -> do + rss_init <- scratch "rss_init" (elemType res_tp) [ry, rx] + (css, ii, jj) <- getThdRedomapRes (rx, ry) (ltid_x, ltid_y) (iii, jjj, redomap_res) + rss <- forLoop ry [rss_init] $ \i [rss_merge] -> do + rss' <- forLoop rx [rss_merge] $ \j [rss_merge'] -> do + prereqAddCode2 (gtid_x, gtid_y) (ii, i, jj, j) (css, redomap_orig_res) + + res_el <- + letSubExp "res_elem" + =<< eIf + ( toExp $ + le64 gtid_y .<. pe64 height_A + .&&. le64 gtid_x .<. pe64 width_B + ) + ( do + addStms code2' + resultBodyM [Var res_nm] + ) + (eBody [eBlank res_tp]) + rss'' <- update "rss" rss_merge' [i, j] res_el + resultBodyM [Var rss''] + resultBodyM [Var rss'] + pure [varRes rss] + let rssss : _ = rssss_list + pure rssss + + let regtile_ret_dims = + map (\(_, sz) -> (sz, se1, se1)) rem_outer_dims + ++ [(height_A, ty, ry), (width_B, tx, rx)] + + -- Add dummy dimensions to tile to reflect the outer dimensions. + epilogue_res' <- + if null rem_outer_dims + then pure epilogue_res + else do + epilogue_t <- lookupType epilogue_res + let (block_dims, rest_dims) = splitAt 2 $ arrayDims epilogue_t + ones = map (const $ intConst Int64 1) rem_outer_dims + new_shape = concat [ones, block_dims, ones, rest_dims] + letExp "res_reshaped" $ BasicOp $ Reshape (map DimNew new_shape) epilogue_res + pure [RegTileReturns mempty regtile_ret_dims epilogue_res'] +mmBlkRegTilingNrm _ _ = pure Nothing + +-- pattern match the properties of the code that we look to +-- tile: a redomap whose two input arrays are each invariant +-- to one of the last two (innermost) parallel dimensions. +matchesBlkRegTile :: + SegSpace -> + Stms GPU -> + Maybe + ( Stms GPU, + (Stm GPU, VName, PrimType, Stm GPU, VName, PrimType), + SubExp, + [Int], + (Lambda GPU, Lambda GPU, SubExp, VName, PrimType) + ) +matchesBlkRegTile seg_space kstms + | -- build the variance table, that records, for + -- each variable name, the variables it depends on + initial_variance <- M.map mempty $ scopeOfSegSpace seg_space, + variance <- varianceInStms initial_variance kstms, + -- check that the code fits the pattern having: + -- some `code1`, followed by one Screma SOAC, followed by some `code2` + (code1, Just screma_stmt, code2) <- matchCodeStreamCode kstms, + Let pat_redomap _ (Op _) <- screma_stmt, + -- checks that the Screma SOAC is actually a redomap and normalizes it + Just (common_dim, arrs, (_, red_lam, red_nes, map_lam)) <- isTileableRedomap screma_stmt, + -- check that exactly two 1D arrays are streamed thorugh redomap, + -- and the result of redomap is one scalar + -- !!!I need to rearrange this whole thing!!! including inp_A and inp_B + length arrs == 2, + [red_ne] <- red_nes, + [map_t1t, map_t2t] <- map paramDec $ lambdaParams map_lam, + [red_t1, _] <- map paramDec $ lambdaParams red_lam, + primType map_t1t && primType map_t2t && primType red_t1, + map_t1_0 <- elemType map_t1t, + map_t2_0 <- elemType map_t2t, + -- checks that the input arrays to redomap are variant to + -- exactly one of the two innermost dimensions of the kernel + Just var_dims <- isInvarTo1of2InnerDims mempty seg_space variance arrs, + -- get the variables on which the first result of redomap depends on + [redomap_orig_res] <- patNames pat_redomap, + Just res_red_var <- M.lookup redomap_orig_res variance, -- variance of the reduce result + + -- we furthermore check that code1 is only formed by + -- 1. statements that slice some globally-declared arrays + -- to produce the input for the redomap, and + -- 2. potentially some statements on which the redomap + -- is independent; these are recorded in `code2''` + Just (code2'', tab_inv_stm) <- + foldl + (processIndirections (namesFromList arrs) res_red_var) + (Just (Seq.empty, M.empty)) + code1, + -- identify load_A, load_B + tmp_stms <- mapMaybe (`M.lookup` tab_inv_stm) arrs, + length tmp_stms == length arrs = + let zip_AB = zip3 tmp_stms arrs [map_t1_0, map_t2_0] + [(load_A, inp_A, map_t1), (load_B, inp_B, map_t2)] = + if var_dims == [0, 1] + then zip_AB + else reverse zip_AB + code2' = code2'' <> code2 + in Just + ( code2', + (load_A, inp_A, map_t1, load_B, inp_B, map_t2), + common_dim, + var_dims, + (map_lam, red_lam, red_ne, redomap_orig_res, elemType red_t1) + ) +matchesBlkRegTile _ _ = Nothing + +-- ceiled division expression +ceilDiv :: MonadBuilder m => SubExp -> SubExp -> m (Exp (Rep m)) +ceilDiv x y = pure $ BasicOp $ BinOp (SDivUp Int64 Unsafe) x y + +mkTileMemSizes :: + SubExp -> + SubExp -> + SubExp -> + Builder + GPU + ( SubExp, + SubExp, + SubExp, + SubExp, + SubExp, + SubExp, + SubExp, + SubExp, + SubExp, + SubExp, + SubExp + ) +mkTileMemSizes height_A width_B common_dim = do + tk_name <- nameFromString . pretty <$> newVName "Tk" + tx_name <- nameFromString . pretty <$> newVName "Tx" + ty_name <- nameFromString . pretty <$> newVName "Ty" + rx_name <- nameFromString . pretty <$> newVName "Rx" + ry_name <- nameFromString . pretty <$> newVName "Ry" + + (ty, ry) <- getParTiles ("Ty", "Ry") (ty_name, ry_name) height_A + (tx, rx) <- getParTiles ("Tx", "Rx") (tx_name, rx_name) width_B + tk <- getSeqTile "Tk" tk_name common_dim ty tx + + tk_div_tx <- letSubExp "tk_div_tx" =<< ceilDiv tk tx + tk_div_ty <- letSubExp "tk_div_ty" =<< ceilDiv tk ty + + tx_rx <- letSubExp "TxRx" =<< toExp (pe64 tx * pe64 rx) + ty_ry <- letSubExp "TyRy" =<< toExp (pe64 ty * pe64 ry) + + let pad_term = sMax64 (pe64 tk) (pe64 ty * pe64 ry) + -- if A not transposed, its shmem should be [ty*ry][tk] + -- we pad to [ty*ry][tk+1] size to minimize bank conflicts + a_loc_sz <- + letSubExp "a_loc_sz" + =<< toExp (pe64 ty * pe64 ry * pe64 tk + pad_term) + -- if B is transposed, its shmem should be [tk][tx*rx] + -- we pad as above, by assuming tx*rx == ty*ry >= tk + -- ToDo: we can decrease the size by checking at this + -- point whether A and B are transposed (or not). + b_loc_sz <- + letSubExp "b_loc_sz" + =<< toExp (pe64 tx * pe64 rx * pe64 tk + pad_term) -- (pe64 tk * pe64 tx * pe64 rx) + pure (rx, ry, tx, ty, tk, tk_div_tx, tk_div_ty, tx_rx, ty_ry, a_loc_sz, b_loc_sz) + +mkNewSegthdLvl :: + SubExp -> + SubExp -> + TPrimExp Int64 VName -> + Builder GPU (SubExp, SubExp, SegLevel) +mkNewSegthdLvl tx ty grid_pexp = do + grid_size <- letSubExp "grid_size" =<< toExp grid_pexp + group_size <- letSubExp "group_size" =<< toExp (pe64 ty * pe64 tx) + let segthd_lvl = SegThread (Count grid_size) (Count group_size) (SegNoVirtFull (SegSeqDims [])) + pure (grid_size, group_size, segthd_lvl) + +mkGidsXYF :: Builder GPU (VName, VName, VName) +mkGidsXYF = do + gid_y <- newVName "gid_y" + gid_x <- newVName "gid_x" + gid_flat <- newVName "gid_flat" + pure (gid_x, gid_y, gid_flat) + +initRegShmem :: + (SubExp, SubExp, SubExp, SubExp, SubExp, SubExp) -> + (PrimType, PrimType, PrimType) -> + SegLevel -> + SubExp -> + Builder GPU (VName, VName, VName) +initRegShmem + (rx, tx, ry, ty, a_loc_sz, b_loc_sz) + (map_t1, map_t2, red_t) + segthd_lvl + red_ne = do + -- initialize register mem with neutral elements. + cssss_list <- segMap2D "cssss" segthd_lvl ResultPrivate (ty, tx) $ \_ -> do + css_init <- scratch "css_init" red_t [ry, rx] + css <- forLoop ry [css_init] $ \i [css_merge] -> do + css' <- forLoop rx [css_merge] $ \j [css_merge'] -> do + css'' <- update "css" css_merge' [i, j] red_ne + resultBodyM [Var css''] + resultBodyM [Var css'] + pure [varRes css] + let [cssss] = cssss_list + -- scratch shared memory + a_loc_init <- scratch "A_loc" map_t1 [a_loc_sz] + b_loc_init <- scratch "B_loc" map_t2 [b_loc_sz] + pure (cssss, a_loc_init, b_loc_init) + +getThdRedomapRes :: + (SubExp, SubExp) -> + (VName, VName) -> + (VName, VName, VName) -> + Builder GPU (VName, VName, VName) +getThdRedomapRes (rx, ry) (ltid_x, ltid_y) (iii, jjj, redomap_res) = do + css <- index "redomap_thd" redomap_res [ltid_y, ltid_x] + ii <- letExp "ii" =<< toExp (le64 iii + le64 ltid_y * pe64 ry) + jj <- letExp "jj" =<< toExp (le64 jjj + le64 ltid_x * pe64 rx) + pure (css, ii, jj) + +prereqAddCode2 :: + (VName, VName) -> + (VName, VName, VName, VName) -> + (VName, VName) -> + Builder GPU () +prereqAddCode2 (gtid_x, gtid_y) (ii, i, jj, j) (css, redomap_orig_res) = do + c <- index "redomap_elm" css [i, j] + cpy_stm <- mkLetNamesM [redomap_orig_res] $ BasicOp $ SubExp $ Var c + addStm cpy_stm + letBindNames [gtid_y] =<< toExp (le64 ii + le64 i) + letBindNames [gtid_x] =<< toExp (le64 jj + le64 j) + +-- | Tries to identify the following pattern: +-- code followed by some Screma followed by more code. +matchCodeStreamCode :: + Stms GPU -> + (Stms GPU, Maybe (Stm GPU), Stms GPU) +matchCodeStreamCode kstms = + let (code1, screma, code2) = + foldl + ( \acc stmt -> + case (acc, stmt) of + ((cd1, Nothing, cd2), Let _ _ (Op (OtherOp Screma {}))) -> + (cd1, Just stmt, cd2) + ((cd1, Nothing, cd2), _) -> + (cd1 ++ [stmt], Nothing, cd2) + ((cd1, Just strm, cd2), _) -> + (cd1, Just strm, cd2 ++ [stmt]) + ) + ([], Nothing, []) + (stmsToList kstms) + in (stmsFromList code1, screma, stmsFromList code2) + +-- | Checks that all streamed arrays are variant to exacly one of +-- the two innermost parallel dimensions, and conversely, for +-- each of the two innermost parallel dimensions, there is at +-- least one streamed array variant to it. The result is the +-- number of the only variant parallel dimension for each array. +isInvarTo1of2InnerDims :: + Names -> + SegSpace -> + VarianceTable -> + [VName] -> + Maybe [Int] +isInvarTo1of2InnerDims branch_variant kspace variance arrs = + let inner_perm0 = map varToOnly1of2InnerDims arrs + inner_perm = catMaybes inner_perm0 + ok1 = elem 0 inner_perm && elem 1 inner_perm + ok2 = length inner_perm0 == length inner_perm + in if ok1 && ok2 then Just inner_perm else Nothing + where + varToOnly1of2InnerDims :: VName -> Maybe Int + varToOnly1of2InnerDims arr = do + (j, _) : (i, _) : _ <- Just $ reverse $ unSegSpace kspace + let variant_to = M.findWithDefault mempty arr variance + branch_invariant = + not $ nameIn j branch_variant || nameIn i branch_variant + if not branch_invariant + then Nothing -- if i or j in branch_variant; return nothing + else + if nameIn i variant_to && not (nameIn j variant_to) + then Just 0 + else + if nameIn j variant_to && not (nameIn i variant_to) + then Just 1 + else Nothing + +processIndirections :: + Names -> -- input arrays to redomap + Names -> -- variables on which the result of redomap depends on. + Maybe (Stms GPU, M.Map VName (Stm GPU)) -> + Stm GPU -> + Maybe (Stms GPU, M.Map VName (Stm GPU)) +processIndirections arrs _ acc stm@(Let patt _ (BasicOp (Index _ _))) + | Just (ss, tab) <- acc, + [p] <- patElems patt, + p_nm <- patElemName p, + nameIn p_nm arrs = + Just (ss, M.insert p_nm stm tab) +processIndirections _ res_red_var acc stm'@(Let patt _ _) + | Just (ss, tab) <- acc, + ps <- patElems patt, + all (\p -> not (nameIn (patElemName p) res_red_var)) ps = + Just (ss Seq.|> stm', tab) + | otherwise = Nothing + +getParTiles :: (String, String) -> (Name, Name) -> SubExp -> Builder GPU (SubExp, SubExp) +getParTiles (t_str, r_str) (t_name, r_name) len_dim = + case len_dim of + Constant (IntValue (Int64Value 8)) -> + pure (se8, se1) + Constant (IntValue (Int64Value 16)) -> + pure (se8, se2) + Constant (IntValue (Int64Value 32)) -> + pure (se8, se4) + _ -> do + t <- letSubExp t_str $ Op $ SizeOp $ GetSize t_name SizeTile + r <- letSubExp r_str $ Op $ SizeOp $ GetSize r_name SizeRegTile + pure (t, r) + +getSeqTile :: String -> Name -> SubExp -> SubExp -> SubExp -> Builder GPU SubExp +getSeqTile tk_str tk_name len_dim ty tx = + case (tx, ty) of + (Constant (IntValue (Int64Value v_x)), Constant (IntValue (Int64Value v_y))) -> + letSubExp tk_str . BasicOp . SubExp . constant $ + case len_dim of + Constant (IntValue (Int64Value v_d)) -> min v_d $ min v_x v_y + _ -> min v_x v_y + _ -> + letSubExp tk_str $ Op $ SizeOp $ GetSize tk_name SizeTile + +---------------------------------------------------------------------------------------------- +--- 3D Tiling (RegTiling for the outermost dimension & Block tiling for the innermost two) --- +---------------------------------------------------------------------------------------------- + +maxRegTile :: Int64 +maxRegTile = 30 + +mkRegTileSe :: Int64 -> SubExp +mkRegTileSe = constant + +variantToDim :: VarianceTable -> VName -> VName -> Bool +variantToDim variance gid_outer nm = + gid_outer == nm || nameIn gid_outer (M.findWithDefault mempty nm variance) + +-- | Checks that all streamed arrays are variant to exacly one of +-- the two innermost parallel dimensions, and conversely, for +-- each of the two innermost parallel dimensions, there is at +-- least one streamed array variant to it. The result is the +-- number of the only variant parallel dimension for each array. +isInvarTo2of3InnerDims :: + Names -> + SegSpace -> + VarianceTable -> + [VName] -> + Maybe [Int] +isInvarTo2of3InnerDims branch_variant kspace variance arrs = + let inner_perm0 = map varToOnly1of3InnerDims arrs + inner_perm = catMaybes inner_perm0 + ok1 = elem 0 inner_perm && elem 1 inner_perm && elem 2 inner_perm + ok2 = length inner_perm0 == length inner_perm + in if ok1 && ok2 then Just inner_perm else Nothing + where + varToOnly1of3InnerDims :: VName -> Maybe Int + varToOnly1of3InnerDims arr = do + (k, _) : (j, _) : (i, _) : _ <- Just $ reverse $ unSegSpace kspace + let variant_to = M.findWithDefault mempty arr variance + branch_invariant = + not $ + nameIn k branch_variant + || nameIn j branch_variant + || nameIn i branch_variant + if not branch_invariant + then Nothing -- if i or j or k in branch_variant; return nothing + else + if nameIn i variant_to && not (nameIn j variant_to || nameIn k variant_to) + then Just 0 + else + if nameIn j variant_to && not (nameIn i variant_to || nameIn k variant_to) + then Just 1 + else + if nameIn k variant_to && not (nameIn i variant_to || nameIn j variant_to) + then Just 2 + else Nothing + +-- | Expects a kernel statement as argument. +-- CONDITIONS for 3D tiling optimization to fire are: +-- 1. a) The kernel body can be broken into +-- scalar-code-1 ++ [Redomap stmt] ++ scalar-code-2. +-- b) The kernels has a per-thread result, and obviously +-- the result is variant to the 3rd dimension +-- (counted from innermost to outermost) +-- 2. For the Redomap: +-- a) the streamed arrays are one dimensional +-- b) each of the array arguments of Redomap are variant +-- to exactly one of the three innermost-parallel dimension +-- of the kernel. This condition can be relaxed by interchanging +-- kernel dimensions whenever possible. +-- 3. For scalar-code-1: +-- a) each of the statements is a slice that produces one of the +-- streamed arrays +-- +-- mmBlkRegTiling :: Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU)) +-- mmBlkRegTiling (Let pat aux (Op (SegOp (SegMap SegThread{} seg_space ts old_kbody)))) +doRegTiling3D :: Stm GPU -> TileM (Maybe (Stms GPU, Stm GPU)) +doRegTiling3D (Let pat aux (Op (SegOp old_kernel))) + | SegMap SegThread {} space kertp (KernelBody () kstms kres) <- old_kernel, + -- build the variance table, that records, for + -- each variable name, the variables it depends on + initial_variance <- M.map mempty $ scopeOfSegSpace space, + variance <- varianceInStms initial_variance kstms, + -- we get the global-thread id for the two inner dimensions, + -- as we are probably going to use it in code generation + (gtid_x, d_Kx) : (gtid_y, d_Ky) : (gtid_z, d_M) : rem_outer_dims_rev <- reverse $ unSegSpace space, + rem_outer_dims <- reverse rem_outer_dims_rev, + -- check that the code fits the pattern having: + -- some `code1`, followed by one Screma SOAC, followed by some `code2` + (code1, Just screma_stmt, code2) <- matchCodeStreamCode kstms, + Let pat_redomap _ (Op _) <- screma_stmt, + -- checks that the Screma SOAC is actually a redomap and normalize it + Just (common_dim, inp_soac_arrs, (_, red_lam, red_nes, map_lam)) <- isTileableRedomap screma_stmt, + not (null red_nes), + -- assuming we have a budget of maxRegTile registers, we distribute + -- that budget across the result of redomap and the kernel result + num_res <- max (length red_nes) (length kres), + reg_tile <- maxRegTile `quot` fromIntegral num_res, + reg_tile_se <- mkRegTileSe reg_tile, + -- check that the element-type of the map and reduce are scalars: + all (primType . paramDec) $ lambdaParams map_lam, + red_res_tps <- map paramDec $ take (length red_nes) $ lambdaParams red_lam, + all primType red_res_tps, + -- checks that the input arrays to redomap are variant to + -- exactly one of the two innermost dimensions of the kernel + Just _ <- isInvarTo2of3InnerDims mempty space variance inp_soac_arrs, + -- get the free variables on which the result of redomap depends on + redomap_orig_res <- patElems pat_redomap, + res_red_var <- -- variance of the reduce result + mconcat $ mapMaybe ((`M.lookup` variance) . patElemName) redomap_orig_res, + mempty /= res_red_var, + -- we furthermore check that code1 is only formed by + -- 1. statements that slice some globally-declared arrays + -- to produce the input for the redomap, and + -- 2. potentially some statements on which the redomap + -- is independent; these are recorded in `code2''` + Just (code2'', arr_tab0) <- + foldl + (processIndirections (namesFromList inp_soac_arrs) res_red_var) + (Just (Seq.empty, M.empty)) + code1, + -- check that code1 contains exacly one slice for each of the input array to redomap + tmp_stms <- mapMaybe (`M.lookup` arr_tab0) inp_soac_arrs, + length tmp_stms == length inp_soac_arrs, + -- code1' <- stmsFromList $ stmsToList code1 \\ stmsToList code2'', + code2' <- code2'' <> code2, + -- we assume the kernel results are variant to the thrid-outer parallel dimension + -- (for sanity sake, they should be) + ker_res_nms <- mapMaybe getResNm kres, + length ker_res_nms == length kres, + all primType kertp, + all (variantToDim variance gtid_z) ker_res_nms = do + -- HERE STARTS THE IMPLEMENTATION: + (new_kernel, host_stms) <- runBuilder $ do + -- host code + -- process the z-variant arrays that need transposition; + -- these "manifest" statements should come before the kernel + (tab_inn, tab_out) <- + foldM + (insertTranspose variance (gtid_z, d_M)) + (M.empty, M.empty) + $ M.toList arr_tab0 + + tx_name <- nameFromString . pretty <$> newVName "Tx" + ty_name <- nameFromString . pretty <$> newVName "Ty" + + tx0 <- letSubExp "Tx" $ Op $ SizeOp $ GetSize tx_name SizeTile + ty0 <- letSubExp "Ty" $ Op $ SizeOp $ GetSize ty_name SizeTile + ty <- limitTile "Ty" ty0 d_Ky + tx <- limitTile "Tx" tx0 d_Kx + let rz = reg_tile_se + + gridDim_x <- letSubExp "gridDim_x" =<< ceilDiv d_Kx tx + gridDim_y <- letSubExp "gridDim_y" =<< ceilDiv d_Ky ty + gridDim_z <- letSubExp "gridDim_z" =<< ceilDiv d_M rz + let gridxyz_pexp = pe64 gridDim_z * pe64 gridDim_y * pe64 gridDim_x + let grid_pexp = product $ gridxyz_pexp : map (pe64 . snd) rem_outer_dims_rev + grid_size <- letSubExp "grid_size_tile3d" =<< toExp grid_pexp + group_size <- letSubExp "group_size_tile3d" =<< toExp (pe64 ty * pe64 tx) + let segthd_lvl = SegThread (Count grid_size) (Count group_size) (SegNoVirtFull (SegSeqDims [])) + + count_shmem <- letSubExp "count_shmem" =<< ceilDiv rz group_size + + gid_x <- newVName "gid_x" + gid_y <- newVName "gid_y" + gid_z <- newVName "gid_z" + gid_flat <- newVName "gid_flat" + + ---- in this binder: outer seggroup ---- + (ret_seggroup, stms_seggroup) <- runBuilder $ do + ii <- letExp "ii" =<< toExp (le64 gid_z * pe64 rz) + jj1 <- letExp "jj1" =<< toExp (le64 gid_y * pe64 ty) + jj2 <- letExp "jj2" =<< toExp (le64 gid_x * pe64 tx) + + -- initialize the register arrays corresponding to the result of redomap; + reg_arr_nms <- segMap2D "res" segthd_lvl ResultPrivate (ty, tx) $ \_ -> + forM (zip red_nes red_res_tps) $ \(red_ne, red_t) -> do + css_init <- scratch "res_init" (elemType red_t) [rz] + css <- forLoop rz [css_init] $ \i [css_merge] -> do + css' <- update "css" css_merge [i] red_ne + resultBodyM [Var css'] + pure $ varRes css + + -- scratch the shared-memory arrays corresponding to the arrays that are + -- input to the redomap and are invariant to the outermost parallel dimension. + loc_arr_nms <- forM (M.toList tab_out) $ \(nm, (ptp, _)) -> + scratch (baseString nm ++ "_loc") ptp [rz] + + prologue_res_list <- + forLoop' common_dim (reg_arr_nms ++ loc_arr_nms) $ + \q var_nms -> do + let reg_arr_merge_nms = take (length red_nes) var_nms + let loc_arr_merge_nms = drop (length red_nes) var_nms + + -- collective copy from global to shared memory + loc_arr_nms' <- + forLoop' count_shmem loc_arr_merge_nms $ \tt loc_arr_merge2_nms -> do + loc_arr_merge2_nms' <- + forM (zip loc_arr_merge2_nms (M.toList tab_out)) $ \(loc_Y_nm, (glb_Y_nm, (ptp_Y, load_Y))) -> do + ltid_flat <- newVName "ltid_flat" + ltid <- newVName "ltid" + let segspace = SegSpace ltid_flat [(ltid, group_size)] + ((res_v, res_i), stms) <- runBuilder $ do + offs <- letExp "offs" =<< toExp (pe64 group_size * le64 tt) + loc_ind <- letExp "loc_ind" =<< toExp (le64 ltid + le64 offs) + letBindNames [gtid_z] =<< toExp (le64 ii + le64 loc_ind) + let glb_ind = gtid_z + y_elm <- + letSubExp "y_elem" + =<< eIf + (toExp $ le64 glb_ind .<. pe64 d_M) + ( do + addStm load_Y + res <- index "Y_elem" glb_Y_nm [q] + resultBodyM [Var res] + ) + (eBody [eBlank $ Prim ptp_Y]) + y_ind <- + letSubExp "y_loc_ind" + =<< eIf + (toExp $ le64 loc_ind .<. pe64 rz) + (toExp loc_ind >>= letTupExp' "loc_fi" >>= resultBodyM) + (eBody [pure $ BasicOp $ SubExp $ intConst Int64 (-1)]) + -- y_tp <- subExpType y_elm + pure (y_elm, y_ind) + + let ret = WriteReturns mempty (Shape [rz]) loc_Y_nm [(Slice [DimFix res_i], res_v)] + let body = KernelBody () stms [ret] + + res_nms <- + letTupExp "Y_glb2loc" <=< renameExp $ + Op $ SegOp $ SegMap segthd_lvl segspace [Prim ptp_Y] body + let res_nm : _ = res_nms + pure res_nm + resultBodyM $ map Var loc_arr_merge2_nms' + + redomap_res <- + segMap2D "redomap_res" segthd_lvl ResultPrivate (ty, tx) $ + \(ltid_y, ltid_x) -> do + letBindNames [gtid_y] =<< toExp (le64 jj1 + le64 ltid_y) + letBindNames [gtid_x] =<< toExp (le64 jj2 + le64 ltid_x) + reg_arr_merge_nms_slc <- forM reg_arr_merge_nms $ \reg_arr_nm -> + index "res_reg_slc" reg_arr_nm [ltid_y, ltid_x] + fmap subExpsRes . letTupExp' "redomap_guarded" + =<< eIf + (toExp $ le64 gtid_y .<. pe64 d_Ky .&&. le64 gtid_x .<. pe64 d_Kx) + ( do + inp_scals_invar_outer <- + forM (M.toList tab_inn) $ \(inp_arr_nm, load_stm) -> do + addStm load_stm + index (baseString inp_arr_nm) inp_arr_nm [q] + -- build the loop of count R whose body is semantically the redomap code + reg_arr_merge_nms' <- + forLoop' rz reg_arr_merge_nms_slc $ \i reg_arr_mm_nms -> do + letBindNames [gtid_z] =<< toExp (le64 ii + le64 i) + resultBodyM =<< letTupExp' "redomap_lam" + =<< eIf + (toExp $ le64 gtid_z .<. pe64 d_M) + ( do + -- read from shared memory + ys <- forM loc_arr_nms' $ \loc_arr_nm -> + index "inp_reg_var2z" loc_arr_nm [i] + cs <- forM reg_arr_mm_nms $ \reg_arr_nm -> + index "res_reg_var2z" reg_arr_nm [i] + -- here we need to put in order the scalar inputs to map: + let tab_scals = + M.fromList $ + zip (map fst $ M.toList tab_out) ys + ++ zip (map fst $ M.toList tab_inn) inp_scals_invar_outer + map_inp_scals <- forM inp_soac_arrs $ \arr_nm -> + case M.lookup arr_nm tab_scals of + Nothing -> error "Impossible case reached in tiling3D\n" + Just nm -> pure nm + map_lam' <- renameLambda map_lam + red_lam' <- renameLambda red_lam + map_res_scals <- eLambda map_lam' (map (eSubExp . Var) map_inp_scals) + red_res <- eLambda red_lam' (map eSubExp (map Var cs ++ map resSubExp map_res_scals)) + css <- forM (zip reg_arr_mm_nms red_res) $ \(reg_arr_nm, c) -> + update (baseString reg_arr_nm) reg_arr_nm [i] (resSubExp c) + resultBodyM $ map Var css + ) + (resultBodyM $ map Var reg_arr_mm_nms) + resultBodyM $ map Var reg_arr_merge_nms' + ) + (resultBodyM $ map Var reg_arr_merge_nms_slc) + resultBodyM $ map Var $ redomap_res ++ loc_arr_nms' + + -- support for non-empty code2' + -- segmap (ltid_y < ty, ltid_x < tx) { + -- for i < rz do + -- res = if (ii+i < d_M && jj1+ltid_y < d_Ky && jj2 + ltid_x < d_Kx) + -- then code2' else dummy + -- final_res[i] = res + let redomap_res = take (length red_nes) prologue_res_list + epilogue_res <- + if length redomap_orig_res == length ker_res_nms + && ker_res_nms == map patElemName redomap_orig_res + then segMap3D "rssss" segthd_lvl ResultPrivate (se1, ty, tx) $ \(_ltid_z, ltid_y, ltid_x) -> forM (zip kertp redomap_res) $ \(res_tp, res) -> do rss_init <- scratch "rss_init" (elemType res_tp) [rz, se1, se1] fmap varRes $ diff --git a/src/Futhark/Optimise/CSE.hs b/src/Futhark/Optimise/CSE.hs --- a/src/Futhark/Optimise/CSE.hs +++ b/src/Futhark/Optimise/CSE.hs @@ -42,6 +42,7 @@ import Futhark.IR import Futhark.IR.Aliases ( Aliases, + consumedInStms, mkStmsAliases, removeFunDefAliases, removeProgAliases, @@ -55,9 +56,6 @@ import Futhark.Pass import Futhark.Transform.Substitute -consumedInStms :: Aliased rep => Stms rep -> Names -consumedInStms = snd . flip mkStmsAliases [] - -- | Perform CSE on every function in a program. -- -- If the boolean argument is false, the pass will not perform CSE on @@ -213,6 +211,13 @@ | patElemName pe `nameIn` consumed = Consume | otherwise = Observe +-- A small amount of normalisation of expressions that otherwise would +-- be different for pointless reasons. +normExp :: Exp lore -> Exp lore +normExp (Apply fname args ret (safety, _, _)) = + Apply fname args ret (safety, mempty, mempty) +normExp e = e + cseInStm :: ASTRep rep => Names -> @@ -221,7 +226,7 @@ CSEM rep a cseInStm consumed (Let pat (StmAux cs attrs edec) e) m = do CSEState (esubsts, nsubsts) cse_arrays <- ask - let e' = substituteNames nsubsts e + let e' = normExp $ substituteNames nsubsts e pat' = substituteNames nsubsts pat if any (bad cse_arrays) $ patElems pat then m [Let pat' (StmAux cs attrs edec) e'] @@ -300,6 +305,8 @@ where cseInOp (GPU.SegOp op) = GPU.SegOp <$> cseInOp op cseInOp (GPU.OtherOp op) = GPU.OtherOp <$> cseInOp op + cseInOp (GPU.GPUBody types body) = + subCSE $ GPU.GPUBody types <$> cseInBody (map (const Observe) types) body cseInOp x = pure x instance diff --git a/src/Futhark/Optimise/Fusion.hs b/src/Futhark/Optimise/Fusion.hs --- a/src/Futhark/Optimise/Fusion.hs +++ b/src/Futhark/Optimise/Fusion.hs @@ -844,10 +844,11 @@ fusionGatherExp _ (Op Futhark.Screma {}) = errorIllegal "screma" fusionGatherExp _ (Op Futhark.Scatter {}) = errorIllegal "write" ------------------------------------ ----- Generic Traversal ---- ------------------------------------ - +fusionGatherExp fres (Op (Futhark.VJP lam _ _)) = + snd <$> fusionGatherLam (mempty, fres) lam +fusionGatherExp fres (Op (Futhark.JVP lam _ _)) = + snd <$> fusionGatherLam (mempty, fres) lam +-- fusionGatherExp fres e = addNamesToInfusible fres $ freeIn e fusionGatherSubExp :: FusedRes -> SubExp -> FusionGM FusedRes diff --git a/src/Futhark/Optimise/GenRedOpt.hs b/src/Futhark/Optimise/GenRedOpt.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/GenRedOpt.hs @@ -0,0 +1,432 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} + +-- | Tries to turn a generalized reduction kernel into +-- a more specialized construct, for example: +-- (a) a map nest with a sequential redomap ripe for tiling +-- (b) a SegRed kernel followed by a smallish accumulation kernel. +-- (c) a histogram (for this we need to track the withAccs) +-- The idea is to identify the first accumulation and +-- to separate the initial kernels into two: +-- 1. the code up to and including the accumulation, +-- which is optimized to turn the accumulation either +-- into a map-reduce composition or a histogram, and +-- 2. the remaining code, which is recursively optimized. +-- Since this is mostly prototyping, when the accumulation +-- can be rewritten as a map-reduce, we sequentialize the +-- map-reduce, as to potentially enable tiling oportunities. +module Futhark.Optimise.GenRedOpt (optimiseGenRed) where + +import Control.Monad.Reader +import Control.Monad.State +import qualified Data.List as L +import qualified Data.Map.Strict as M +import Data.Maybe +import Futhark.Builder +import Futhark.IR.GPU +import Futhark.Optimise.TileLoops.Shared +import Futhark.Pass +import Futhark.Tools +import Futhark.Transform.Rename + +type GenRedM = ReaderT (Scope GPU) (State VNameSource) + +-- | The pass definition. +optimiseGenRed :: Pass GPU GPU +optimiseGenRed = + Pass "optimise generalized reductions" "Specializes generalized reductions into map-reductions or histograms" $ + intraproceduralTransformation onStms + where + onStms scope stms = + modifyNameSource $ + runState $ + runReaderT (optimiseStms (M.empty, M.empty) stms) scope + +optimiseBody :: Env -> Body GPU -> GenRedM (Body GPU) +optimiseBody env (Body () stms res) = + Body () <$> optimiseStms env stms <*> pure res + +optimiseStms :: Env -> Stms GPU -> GenRedM (Stms GPU) +optimiseStms env stms = + localScope (scopeOf stms) $ do + (_, stms') <- foldM foldfun (env, mempty) $ stmsToList stms + pure stms' + where + foldfun :: (Env, Stms GPU) -> Stm GPU -> GenRedM (Env, Stms GPU) + foldfun (e, ss) s = do + (e', s') <- optimiseStm e s + pure (e', ss <> s') + +optimiseStm :: Env -> Stm GPU -> GenRedM (Env, Stms GPU) +optimiseStm env stm@(Let _ _ (Op (SegOp (SegMap SegThread {} _ _ _)))) = do + res_genred_opt <- genRedOpts env stm + let stms' = + case res_genred_opt of + Just stms -> stms + Nothing -> oneStm stm + pure (env, stms') +optimiseStm env (Let pat aux e) = do + env' <- changeEnv env (head $ patNames pat) e + e' <- mapExpM (optimise env') e + pure (env', oneStm $ Let pat aux e') + where + optimise env' = identityMapper {mapOnBody = \scope -> localScope scope . optimiseBody env'} + +------------------------ + +genRedOpts :: Env -> Stm GPU -> GenRedM (Maybe (Stms GPU)) +genRedOpts env ker = do + res_tile <- genRed2Tile2d env ker + case res_tile of + Nothing -> do + res_sgrd <- genRed2SegRed env ker + helperGenRed res_sgrd + _ -> helperGenRed res_tile + where + helperGenRed Nothing = pure Nothing + helperGenRed (Just (stms_before, ker_snd)) = do + mb_stms_after <- genRedOpts env ker_snd + case mb_stms_after of + Just stms_after -> pure $ Just $ stms_before <> stms_after + Nothing -> pure $ Just $ stms_before <> oneStm ker_snd + +se1 :: SubExp +se1 = intConst Int64 1 + +genRed2Tile2d :: Env -> Stm GPU -> GenRedM (Maybe (Stms GPU, Stm GPU)) +genRed2Tile2d env kerstm@(Let pat_ker aux (Op (SegOp (SegMap seg_thd seg_space kres_tps old_kbody)))) + | (SegThread _ seg_group_size _novirt) <- seg_thd, + -- novirt == SegNoVirtFull || novirt == SegNoVirt, + KernelBody () kstms kres <- old_kbody, + Just (css, r_ses) <- allGoodReturns kres, + null css, + -- build the variance table, that records, for + -- each variable name, the variables it depends on + initial_variance <- M.map mempty $ scopeOfSegSpace seg_space, + variance <- varianceInStms initial_variance kstms, + -- check that the code fits the pattern having: + -- some `code1`, followed by one accumulation, followed by some `code2` + -- UpdateAcc VName [SubExp] [SubExp] + (code1, Just accum_stmt, code2) <- matchCodeAccumCode kstms, + Let pat_accum _aux_acc (BasicOp (UpdateAcc acc_nm acc_inds acc_vals)) <- accum_stmt, + [pat_acc_nm] <- patNames pat_accum, + -- check that the `acc_inds` are invariant to at least one + -- parallel kernel dimensions, and return the innermost such one: + Just (invar_gid, gid_ind) <- isInvarToParDim mempty seg_space variance acc_inds, + gid_dims_new_0 <- filter (\x -> invar_gid /= fst x) (unSegSpace seg_space), + -- reorder the variant dimensions such that inner(most) accum-indices + -- correspond to inner(most) parallel dimensions, so that the babysitter + -- does not introduce transpositions + -- gid_dims_new <- gid_dims_new_0, + gid_dims_new <- reorderParDims variance acc_inds gid_dims_new_0, + -- check that all global-memory accesses in `code1` on which + -- `accum_stmt` depends on are invariant to at least one of + -- the remaining parallel dimensions (i.e., excluding `invar_gid`) + all (isTileable invar_gid gid_dims_new variance pat_acc_nm) (stmsToList code1), + -- need to establish a cost model for the stms that would now + -- be redundantly executed by the two kernels. If any recurence + -- is redundant than it is a no go. Otherwise we need to look at + -- memory accesses: if more than two are re-executed, then we + -- should abort. + cost <- costRedundantExecution variance pat_acc_nm r_ses kstms, + maxCost cost (Small 2) == Small 2 = do + -- 1. create the first kernel + acc_tp <- lookupType acc_nm + let inv_dim_len = segSpaceDims seg_space !! gid_ind + -- 1.1. get the accumulation operator + ((redop0, neutral), el_tps) = getAccLambda acc_tp + redop <- renameLambda redop0 + let red = + Reduce + { redComm = Commutative, + redLambda = redop, + redNeutral = neutral + } + -- 1.2. build the sequential map-reduce screma + code1' = + stmsFromList $ + filter (dependsOnAcc pat_acc_nm variance) $ + stmsToList code1 + (code1'', code1_tr_host) <- transposeFVs (freeIn kerstm) variance invar_gid code1' + let map_lam_body = mkBody code1'' $ map (SubExpRes (Certs [])) acc_vals + map_lam0 = Lambda [Param mempty invar_gid (Prim int64)] map_lam_body el_tps + map_lam <- renameLambda map_lam0 + (k1_res, ker1_stms) <- runBuilderT' $ do + iota <- letExp "iota" $ BasicOp $ Iota inv_dim_len (intConst Int64 0) (intConst Int64 1) Int64 + let op_exp = Op (OtherOp (Screma inv_dim_len [iota] (ScremaForm [] [red] map_lam))) + res_redmap <- letTupExp "res_mapred" op_exp + letSubExp (baseString pat_acc_nm ++ "_big_update") $ + BasicOp (UpdateAcc acc_nm acc_inds $ map Var res_redmap) + + -- 1.3. build the kernel expression and rename it! + gid_flat_1 <- newVName "gid_flat" + let space1 = SegSpace gid_flat_1 gid_dims_new + + (grid_size, host_stms1) <- runBuilder $ do + let grid_pexp = foldl (\x d -> x * pe64 d) (pe64 se1) $ map snd gid_dims_new + dim_prod <- letSubExp "dim_prod" =<< toExp grid_pexp + letSubExp "grid_size" =<< ceilDiv dim_prod (unCount seg_group_size) + let level1 = SegThread (Count grid_size) seg_group_size (SegNoVirtFull (SegSeqDims [])) -- novirt ? + kbody1 = KernelBody () ker1_stms [Returns ResultMaySimplify (Certs []) k1_res] + + -- is it OK here to use the "aux" from the parrent kernel? + ker_exp <- renameExp $ Op (SegOp (SegMap level1 space1 [acc_tp] kbody1)) + let ker1 = Let pat_accum aux ker_exp + + -- 2 build the second kernel + let ker2_body = old_kbody {kernelBodyStms = code1 <> code2} + ker2_exp <- renameExp $ Op (SegOp (SegMap seg_thd seg_space kres_tps ker2_body)) + let ker2 = Let pat_ker aux ker2_exp + pure $ + Just (code1_tr_host <> host_stms1 <> oneStm ker1, ker2) + where + isIndVarToParDim _ (Constant _) _ = False + isIndVarToParDim variance (Var acc_ind) par_dim = + acc_ind == fst par_dim + || nameIn (fst par_dim) (M.findWithDefault mempty acc_ind variance) + foldfunReorder variance (unused_dims, inner_dims) acc_ind = + case L.findIndex (isIndVarToParDim variance acc_ind) unused_dims of + Nothing -> (unused_dims, inner_dims) + Just i -> + ( take i unused_dims ++ drop (i + 1) unused_dims, + (unused_dims !! i) : inner_dims + ) + reorderParDims variance acc_inds gid_dims_new_0 = + let (invar_dims, inner_dims) = + foldl + (foldfunReorder variance) + (gid_dims_new_0, []) + (reverse acc_inds) + in invar_dims ++ inner_dims + -- + ceilDiv x y = pure $ BasicOp $ BinOp (SDivUp Int64 Unsafe) x y + getAccLambda acc_tp = + case acc_tp of + (Acc tp_id _shp el_tps _) -> + case M.lookup tp_id (fst env) of + Just lam -> (lam, el_tps) + _ -> error $ "Lookup in environment failed! " ++ pretty tp_id ++ " env: " ++ pretty (fst env) + _ -> error "Illegal accumulator type!" + -- is a subexp invariant to a gid of a parallel dimension? + isSeInvar2 variance gid (Var x) = + let x_deps = M.findWithDefault mempty x variance + in gid /= x && not (nameIn gid x_deps) + isSeInvar2 _ _ _ = True + -- is a DimIndex invar to a gid of a parallel dimension? + isDimIdxInvar2 variance gid (DimFix d) = + isSeInvar2 variance gid d + isDimIdxInvar2 variance gid (DimSlice d1 d2 d3) = + all (isSeInvar2 variance gid) [d1, d2, d3] + -- is an entire slice invariant to at least one gid of a parallel dimension + isSliceInvar2 variance slc = + any (\gid -> all (isDimIdxInvar2 variance gid) (unSlice slc)) + -- are all statements that touch memory invariant to at least one parallel dimension? + isTileable :: VName -> [(VName, SubExp)] -> VarianceTable -> VName -> Stm GPU -> Bool + isTileable seq_gid gid_dims variance acc_nm (Let (Pat [pel]) _ (BasicOp (Index _ slc))) + | acc_deps <- M.findWithDefault mempty acc_nm variance, + nameIn (patElemName pel) acc_deps = + let invar_par = isSliceInvar2 variance slc (map fst gid_dims) + invar_seq = isSliceInvar2 variance slc [seq_gid] + in invar_par || invar_seq + -- this relies on the cost model, that currently accepts only + -- global-memory reads, and for example rejects in-place updates + -- or loops inside the code that is transformed in a redomap. + isTileable _ _ _ _ _ = True + -- does the to-be-reduced accumulator depends on this statement? + dependsOnAcc pat_acc_nm variance (Let pat _ _) = + let acc_deps = M.findWithDefault mempty pat_acc_nm variance + in any (`nameIn` acc_deps) $ patNames pat +genRed2Tile2d _ _ = + pure Nothing + +genRed2SegRed :: Env -> Stm GPU -> GenRedM (Maybe (Stms GPU, Stm GPU)) +genRed2SegRed _ _ = + pure Nothing + +transposeFVs :: + Names -> + VarianceTable -> + VName -> + Stms GPU -> + GenRedM (Stms GPU, Stms GPU) +transposeFVs fvs variance gid stms = do + (tab, stms') <- foldM foldfun (M.empty, mempty) $ stmsToList stms + let stms_host = M.foldr (\(_, _, s) ss -> ss <> s) mempty tab + pure (stms', stms_host) + where + foldfun (tab, all_stms) stm = do + (tab', stm') <- transposeFV (tab, stm) + pure (tab', all_stms <> oneStm stm') + -- ToDo: currently handles only 2-dim arrays, please generalize + transposeFV (tab, Let pat aux (BasicOp (Index arr slc))) + | dims <- unSlice slc, + all isFixDim dims, + nameIn arr fvs, + iis <- L.findIndices depOnGid dims, + [ii] <- iis, + -- generalize below: treat any rearange and add to tab if not there. + Nothing <- M.lookup arr tab, + ii /= length dims - 1, + perm <- [0 .. ii - 1] ++ [ii + 1 .. length dims - 1] ++ [ii] = do + (arr_tr, stms_tr) <- runBuilderT' $ do + arr' <- letExp (baseString arr ++ "_trsp") $ BasicOp $ Rearrange perm arr -- Manifest [1,0] arr + letExp (baseString arr' ++ "_opaque") $ BasicOp $ Opaque OpaqueNil $ Var arr' + let tab' = M.insert arr (perm, arr_tr, stms_tr) tab + slc' = Slice $ map (dims !!) perm + stm' = Let pat aux $ BasicOp $ Index arr_tr slc' + pure (tab', stm') + where + isFixDim DimFix {} = True + isFixDim _ = False + depOnGid (DimFix (Var nm)) = + gid == nm || nameIn gid (M.findWithDefault mempty nm variance) + depOnGid _ = False + transposeFV r = pure r + +-- | Tries to identify the following pattern: +-- code followed by some UpdateAcc-statement +-- followed by more code. +matchCodeAccumCode :: + Stms GPU -> + (Stms GPU, Maybe (Stm GPU), Stms GPU) +matchCodeAccumCode kstms = + let (code1, screma, code2) = + foldl + ( \acc stmt -> + case (acc, stmt) of + ((cd1, Nothing, cd2), Let _ _ (BasicOp UpdateAcc {})) -> + (cd1, Just stmt, cd2) + ((cd1, Nothing, cd2), _) -> + (cd1 ++ [stmt], Nothing, cd2) + ((cd1, Just strm, cd2), _) -> + (cd1, Just strm, cd2 ++ [stmt]) + ) + ([], Nothing, []) + (stmsToList kstms) + in (stmsFromList code1, screma, stmsFromList code2) + +-- | Checks that there exist a parallel dimension (among @kids@), +-- to which all the indices (@acc_inds@) are invariant to. +-- It returns the innermost such parallel dimension, as a tuple +-- of the pardim gid ('VName') and its index ('Int') in the +-- parallel space. +isInvarToParDim :: + Names -> + SegSpace -> + VarianceTable -> + [SubExp] -> + Maybe (VName, Int) +isInvarToParDim branch_variant kspace variance acc_inds = + let ker_gids = map fst $ unSegSpace kspace + branch_invariant = not $ any (`nameIn` branch_variant) ker_gids + allvar2 = allvariant2 acc_inds ker_gids + last_invar_dim = + foldl (lastNotIn allvar2) Nothing $ + zip ker_gids [0 .. length ker_gids - 1] + in if branch_invariant + then last_invar_dim + else Nothing + where + variant2 (Var ind) kids = + let variant_to = + M.findWithDefault mempty ind variance + <> (if ind `elem` kids then oneName ind else mempty) + in filter (`nameIn` variant_to) kids + variant2 _ _ = [] + allvariant2 ind_ses kids = + namesFromList $ concatMap (`variant2` kids) ind_ses + lastNotIn allvar2 acc (kid, k) = + if nameIn kid allvar2 then acc else Just (kid, k) + +allGoodReturns :: [KernelResult] -> Maybe ([VName], [SubExp]) +allGoodReturns kres + | all goodReturn kres = do + Just $ foldl addCertAndRes ([], []) kres + where + goodReturn (Returns ResultMaySimplify _ _) = True + goodReturn _ = False + addCertAndRes (cs, rs) (Returns ResultMaySimplify c r_se) = + (cs ++ unCerts c, rs ++ [r_se]) + addCertAndRes _ _ = + error "Impossible case reached in GenRedOpt.hs, function allGoodReturns!" +allGoodReturns _ = Nothing + +-------------------------- +--- Cost Model Helpers --- +-------------------------- + +costRedundantExecution :: + VarianceTable -> + VName -> + [SubExp] -> + Stms GPU -> + Cost +costRedundantExecution variance pat_acc_nm r_ses kstms = + let acc_deps = M.findWithDefault mempty pat_acc_nm variance + vartab_cut_acc = varianceInStmsWithout (oneName pat_acc_nm) mempty kstms + res_deps = mconcat $ map (findDeps vartab_cut_acc) $ mapMaybe se2nm r_ses + common_deps = namesIntersection res_deps acc_deps + in foldl (addCostOfStmt common_deps) (Small 0) kstms + where + se2nm (Var nm) = Just nm + se2nm _ = Nothing + findDeps vartab nm = M.findWithDefault mempty nm vartab + addCostOfStmt common_deps cur_cost stm = + let pat_nms = patNames $ stmPat stm + in if namesIntersect (namesFromList pat_nms) common_deps + then addCosts cur_cost $ costRedundantStmt stm + else cur_cost + varianceInStmsWithout :: Names -> VarianceTable -> Stms GPU -> VarianceTable + varianceInStmsWithout nms = L.foldl' (varianceInStmWithout nms) + varianceInStmWithout cuts vartab stm = + let pat_nms = patNames $ stmPat stm + in if namesIntersect (namesFromList pat_nms) cuts + then vartab + else L.foldl' add vartab pat_nms + where + add variance' v = M.insert v binding_variance variance' + look variance' v = oneName v <> M.findWithDefault mempty v variance' + binding_variance = mconcat $ map (look vartab) $ namesToList (freeIn stm) + +data Cost = Small Int | Big | Break + deriving (Eq) + +addCosts :: Cost -> Cost -> Cost +addCosts Break _ = Break +addCosts _ Break = Break +addCosts Big _ = Big +addCosts _ Big = Big +addCosts (Small c1) (Small c2) = Small (c1 + c2) + +maxCost :: Cost -> Cost -> Cost +maxCost (Small c1) (Small c2) = Small (max c1 c2) +maxCost c1 c2 = addCosts c1 c2 + +costBody :: Body GPU -> Cost +costBody bdy = + foldl addCosts (Small 0) $ + map costRedundantStmt $ stmsToList $ bodyStms bdy + +costRedundantStmt :: Stm GPU -> Cost +costRedundantStmt (Let _ _ (Op _)) = Big +costRedundantStmt (Let _ _ DoLoop {}) = Big +costRedundantStmt (Let _ _ Apply {}) = Big +costRedundantStmt (Let _ _ WithAcc {}) = Big +costRedundantStmt (Let _ _ (If _cond b_then b_else _)) = + maxCost (costBody b_then) (costBody b_else) +costRedundantStmt (Let _ _ (BasicOp (ArrayLit _ Array {}))) = Big +costRedundantStmt (Let _ _ (BasicOp (ArrayLit _ _))) = Small 1 +costRedundantStmt (Let _ _ (BasicOp (Index _ slc))) = + if all isFixDim (unSlice slc) then Small 1 else Small 0 + where + isFixDim DimFix {} = True + isFixDim _ = False +costRedundantStmt (Let _ _ (BasicOp FlatIndex {})) = Small 0 +costRedundantStmt (Let _ _ (BasicOp Update {})) = Break +costRedundantStmt (Let _ _ (BasicOp FlatUpdate {})) = Break +costRedundantStmt (Let _ _ (BasicOp Concat {})) = Big +costRedundantStmt (Let _ _ (BasicOp Copy {})) = Big +costRedundantStmt (Let _ _ (BasicOp Manifest {})) = Big +costRedundantStmt (Let _ _ (BasicOp Replicate {})) = Big +costRedundantStmt (Let _ _ (BasicOp UpdateAcc {})) = Break +costRedundantStmt (Let _ _ (BasicOp _)) = Small 0 diff --git a/src/Futhark/Optimise/HistAccs.hs b/src/Futhark/Optimise/HistAccs.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/HistAccs.hs @@ -0,0 +1,163 @@ +{-# LANGUAGE TypeFamilies #-} + +-- | Turn certain uses of accumulators into SegHists. +module Futhark.Optimise.HistAccs (histAccsGPU) where + +import Control.Monad.Reader +import Control.Monad.State +import qualified Data.Map.Strict as M +import Futhark.IR.GPU +import Futhark.MonadFreshNames +import Futhark.Pass +import Futhark.Tools +import Futhark.Transform.Rename +import Prelude hiding (quot) + +-- | A mapping from accumulator variables to their source. +type Accs rep = M.Map VName (WithAccInput rep) + +type OptM = ReaderT (Scope GPU) (State VNameSource) + +optimiseBody :: Accs GPU -> Body GPU -> OptM (Body GPU) +optimiseBody accs body = mkBody <$> optimiseStms accs (bodyStms body) <*> pure (bodyResult body) + +optimiseExp :: Accs GPU -> Exp GPU -> OptM (Exp GPU) +optimiseExp accs = mapExpM mapper + where + mapper = + identityMapper + { mapOnBody = \scope body -> localScope scope $ optimiseBody accs body + } + +extractUpdate :: + Accs rep -> + VName -> + Stms rep -> + Maybe ((WithAccInput rep, VName, [SubExp], [SubExp]), Stms rep) +extractUpdate accs v stms = do + (stm, stms') <- stmsHead stms + case stm of + Let (Pat [PatElem pe_v _]) _ (BasicOp (UpdateAcc acc is vs)) + | pe_v == v -> do + acc_input <- M.lookup acc accs + Just ((acc_input, acc, is, vs), stms') + _ -> do + (x, stms'') <- extractUpdate accs v stms' + pure (x, oneStm stm <> stms'') + +mkHistBody :: Accs GPU -> KernelBody GPU -> Maybe (KernelBody GPU, WithAccInput GPU, VName) +mkHistBody accs (KernelBody () stms [Returns rm cs (Var v)]) = do + ((acc_input, acc, is, vs), stms') <- extractUpdate accs v stms + pure + ( KernelBody () stms' $ map (Returns rm cs) is ++ map (Returns rm cs) vs, + acc_input, + acc + ) +mkHistBody _ _ = Nothing + +withAccLamToHistLam :: MonadFreshNames m => Shape -> Lambda GPU -> m (Lambda GPU) +withAccLamToHistLam shape lam = + renameLambda $ lam {lambdaParams = drop (shapeRank shape) (lambdaParams lam)} + +addArrsToAcc :: + (MonadBuilder m, Rep m ~ GPU) => + SegLevel -> + Shape -> + [VName] -> + VName -> + m (Exp GPU) +addArrsToAcc lvl shape arrs acc = do + flat <- newVName "phys_tid" + gtids <- replicateM (shapeRank shape) (newVName "gtid") + let space = SegSpace flat $ zip gtids $ shapeDims shape + + (acc', stms) <- localScope (scopeOfSegSpace space) . collectStms $ do + vs <- forM arrs $ \arr -> do + arr_t <- lookupType arr + letSubExp (baseString arr <> "_elem") $ + BasicOp $ Index arr $ fullSlice arr_t $ map (DimFix . Var) gtids + letExp (baseString acc <> "_upd") $ + BasicOp $ UpdateAcc acc (map Var gtids) vs + + acc_t <- lookupType acc + pure . Op . SegOp . SegMap lvl space [acc_t] $ + KernelBody () stms [Returns ResultMaySimplify mempty (Var acc')] + +flatKernelBody :: + MonadBuilder m => + SegSpace -> + KernelBody (Rep m) -> + m (SegSpace, KernelBody (Rep m)) +flatKernelBody space kbody = do + gtid <- newVName "gtid" + dims_prod <- + letSubExp "dims_prod" + =<< foldBinOp (Mul Int64 OverflowUndef) (intConst Int64 1) (segSpaceDims space) + + let space' = SegSpace (segFlat space) [(gtid, dims_prod)] + + kbody_stms <- localScope (scopeOfSegSpace space') . collectStms_ $ do + let new_inds = + unflattenIndex (map pe64 (segSpaceDims space)) (pe64 $ Var gtid) + zipWithM_ letBindNames (map (pure . fst) (unSegSpace space)) + =<< mapM toExp new_inds + addStms $ kernelBodyStms kbody + + pure (space', kbody {kernelBodyStms = kbody_stms}) + +optimiseStm :: Accs GPU -> Stm GPU -> OptM (Stms GPU) +-- TODO: this is very restricted currently, but shows the idea. +optimiseStm accs (Let pat aux (WithAcc inputs lam)) = do + localScope (scopeOfLParams (lambdaParams lam)) $ do + body' <- optimiseBody accs' $ lambdaBody lam + let lam' = lam {lambdaBody = body'} + pure $ oneStm $ Let pat aux $ WithAcc inputs lam' + where + acc_names = map paramName $ drop (length inputs) $ lambdaParams lam + accs' = M.fromList (zip acc_names inputs) <> accs +optimiseStm accs (Let pat aux (Op (SegOp (SegMap lvl space _ kbody)))) + | accs /= mempty, + Just (kbody', (acc_shape, _, Just (acc_lam, acc_nes)), acc) <- + mkHistBody accs kbody, + all primType $ lambdaReturnType acc_lam = runBuilder_ $ do + hist_dests <- forM acc_nes $ \ne -> + letExp "hist_dest" $ BasicOp $ Replicate acc_shape ne + + acc_lam' <- withAccLamToHistLam acc_shape acc_lam + + let ts' = + replicate (shapeRank acc_shape) (Prim int64) + ++ lambdaReturnType acc_lam + histop = + HistOp + { histShape = acc_shape, + histRaceFactor = intConst Int64 1, + histDest = hist_dests, + histNeutral = acc_nes, + histOpShape = mempty, + histOp = acc_lam' + } + + (space', kbody'') <- flatKernelBody space kbody' + + hist_dest_upd <- + letTupExp "hist_dest_upd" $ Op $ SegOp $ SegHist lvl space' [histop] ts' kbody'' + + addStm . Let pat aux =<< addArrsToAcc lvl acc_shape hist_dest_upd acc +optimiseStm accs (Let pat aux e) = + oneStm . Let pat aux <$> optimiseExp accs e + +optimiseStms :: Accs GPU -> Stms GPU -> OptM (Stms GPU) +optimiseStms accs stms = + localScope (scopeOf stms) $ + mconcat <$> mapM (optimiseStm accs) (stmsToList stms) + +-- | The pass for GPU kernels. +histAccsGPU :: Pass GPU GPU +histAccsGPU = + Pass "hist accs" "Turn certain accumulations into histograms" $ + intraproceduralTransformation onStms + where + onStms scope stms = + modifyNameSource . runState $ + runReaderT (optimiseStms mempty stms) scope diff --git a/src/Futhark/Optimise/MergeGPUBodies.hs b/src/Futhark/Optimise/MergeGPUBodies.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/MergeGPUBodies.hs @@ -0,0 +1,634 @@ +-- | +-- This module implements an optimization pass that merges 'GPUBody' kernels to +-- eliminate memory transactions and reduce the number of kernel launches. +-- This is useful because the "Futhark.Optimise.ReduceDeviceSyncs" pass introduces +-- 'GPUBody' kernels that only execute single statements. +-- +-- To merge as many 'GPUBody' kernels as possible, this pass reorders statements +-- with the goal of bringing as many 'GPUBody' statements next to each other in +-- a sequence. Such sequence can then trivially be merged. +module Futhark.Optimise.MergeGPUBodies (mergeGPUBodies) where + +import Control.Monad +import Control.Monad.Trans.Class +import Control.Monad.Trans.State.Strict hiding (State) +import Data.Foldable +import qualified Data.IntMap as IM +import Data.IntSet ((\\)) +import qualified Data.IntSet as IS +import qualified Data.Map as M +import Data.Maybe (fromMaybe) +import Data.Sequence ((|>)) +import qualified Data.Sequence as SQ +import Futhark.Analysis.Alias +import Futhark.Construct (sliceDim) +import Futhark.Error +import Futhark.IR.Aliases +import Futhark.IR.GPU +import Futhark.MonadFreshNames hiding (newName) +import Futhark.Pass + +-- | An optimization pass that reorders and merges 'GPUBody' statements to +-- eliminate memory transactions and reduce the number of kernel launches. +mergeGPUBodies :: Pass GPU GPU +mergeGPUBodies = + Pass + "merge GPU bodies" + "Reorder and merge GPUBody constructs to reduce kernels executions." + $ intraproceduralTransformationWithConsts onStms onFunDef . aliasAnalysis + where + onFunDef _ (FunDef entry attrs name types params body) = + FunDef entry attrs name types params . fst <$> transformBody mempty body + onStms stms = + fst <$> transformStms mempty stms + +-------------------------------------------------------------------------------- +-- COMMON - TYPES -- +-------------------------------------------------------------------------------- + +-- | A set of 'VName' tags that denote all variables that some group of +-- statements depend upon. Those must be computed before the group statements. +type Dependencies = IS.IntSet + +-- | A set of 'VName' tags that denote all variables that some group of +-- statements binds. +type Bindings = IS.IntSet + +-- | A set of 'VName' tags that denote the root aliases of all arrays that some +-- statement consumes. +type Consumption = IS.IntSet + +-------------------------------------------------------------------------------- +-- COMMON - HELPERS -- +-------------------------------------------------------------------------------- + +-- | All free variables of a construct as 'Dependencies'. +depsOf :: FreeIn a => a -> Dependencies +depsOf = namesToSet . freeIn + +-- | Convert 'Names' to an integer set of name tags. +namesToSet :: Names -> IS.IntSet +namesToSet = IS.fromList . map baseTag . namesToList + +-------------------------------------------------------------------------------- +-- AD HOC OPTIMIZATION -- +-------------------------------------------------------------------------------- + +-- | Optimize a lambda and determine its dependencies. +transformLambda :: + AliasTable -> + Lambda (Aliases GPU) -> + PassM (Lambda GPU, Dependencies) +transformLambda aliases (Lambda params body types) = do + (body', deps) <- transformBody aliases body + pure (Lambda params body' types, deps) + +-- | Optimize a body and determine its dependencies. +transformBody :: + AliasTable -> + Body (Aliases GPU) -> + PassM (Body GPU, Dependencies) +transformBody aliases (Body _ stms res) = do + grp <- evalStateT (foldM_ reorderStm aliases stms >> collapse) initialState + + let stms' = groupStms grp + let deps = (groupDependencies grp <> depsOf res) \\ groupBindings grp + + pure (Body () stms' res, deps) + +-- | Optimize a sequence of statements and determine their dependencies. +transformStms :: + AliasTable -> + Stms (Aliases GPU) -> + PassM (Stms GPU, Dependencies) +transformStms aliases stms = do + (Body _ stms' _, deps) <- transformBody aliases (Body mempty stms []) + pure (stms', deps) + +-- | Optimizes and reorders a single statement within a sequence while tracking +-- the declaration, observation, and consumption of its dependencies. +-- This creates sequences of GPUBody statements that can be merged into single +-- kernels. +reorderStm :: AliasTable -> Stm (Aliases GPU) -> ReorderM AliasTable +reorderStm aliases (Let pat (StmAux cs attrs _) e) = do + (e', deps) <- lift (transformExp aliases e) + let pat' = removePatAliases pat + let stm' = Let pat' (StmAux cs attrs ()) e' + let pes' = patElems pat' + + -- Array aliases can be seen as a directed graph where vertices are arrays + -- (or the names that bind them) and an edge x -> y denotes that x aliases y. + -- The root aliases of some array A is then the set of arrays that can be + -- reached from A in graph and which have no edges themselves. + -- + -- All arrays that share a root alias are considered aliases of each other + -- and will be consumed if either of them is consumed. + -- When reordering statements we must ensure that no statement that consumes + -- an array is moved before any statement that observes one of its aliases. + -- + -- That is to move statement X before statement Y the set of root aliases of + -- arrays consumed by X must not overlap with the root aliases of arrays + -- observed by Y. + -- + -- We consider the root aliases of Y's observed arrays as part of Y's + -- dependencies and simply say that the root aliases of arrays consumed by X + -- must not overlap those. + -- + -- To move X before Y then the dependencies of X must also not overlap with + -- the variables bound by Y. + + let observed = namesToSet $ rootAliasesOf (fold $ expAliases e) aliases + let consumed = namesToSet $ rootAliasesOf (consumedInExp e) aliases + let usage = + Usage + { usageBindings = IS.fromList $ map (baseTag . patElemName) pes', + usageDependencies = observed <> deps <> depsOf pat' <> depsOf cs + } + + case e' of + Op GPUBody {} -> + moveGPUBody stm' usage consumed + _ -> + moveOther stm' usage consumed + + pure $ foldl recordAliases aliases (patElems pat) + where + rootAliasesOf names atable = + let look n = M.findWithDefault (oneName n) n atable + in foldMap look (namesToList names) + + recordAliases atable pe + | aliasesOf pe == mempty = + atable + | otherwise = + let root_aliases = rootAliasesOf (aliasesOf pe) atable + in M.insert (patElemName pe) root_aliases atable + +-- | Optimize a single expression and determine its dependencies. +transformExp :: + AliasTable -> + Exp (Aliases GPU) -> + PassM (Exp GPU, Dependencies) +transformExp aliases e = + case e of + BasicOp {} -> pure (removeExpAliases e, depsOf e) + Apply {} -> pure (removeExpAliases e, depsOf e) + If c tbody fbody dec -> do + (tbody', t_deps) <- transformBody aliases tbody + (fbody', f_deps) <- transformBody aliases fbody + let deps = depsOf c <> t_deps <> f_deps <> depsOf dec + pure (If c tbody' fbody' dec, deps) + DoLoop merge lform body -> do + -- What merge and lform aliases outside the loop is irrelevant as those + -- cannot be consumed within the loop. + (body', body_deps) <- transformBody aliases body + let (params, args) = unzip merge + let deps = body_deps <> depsOf params <> depsOf args <> depsOf lform + + let scope = scopeOf lform <> scopeOfFParams params + let bound = IS.fromList $ map baseTag (M.keys scope) + let deps' = deps \\ bound + + let dummy = DoLoop merge lform (Body (bodyDec body) SQ.empty []) + let DoLoop merge' lform' _ = removeExpAliases dummy + + pure (DoLoop merge' lform' body', deps') + WithAcc inputs lambda -> do + accs <- mapM (transformWithAccInput aliases) inputs + let (inputs', input_deps) = unzip accs + -- The lambda parameters are all unique and thus have no aliases. + (lambda', deps) <- transformLambda aliases lambda + pure (WithAcc inputs' lambda', deps <> fold input_deps) + Op {} -> + -- A GPUBody cannot be nested within other HostOp constructs. + pure (removeExpAliases e, depsOf e) + +-- | Optimize a single WithAcc input and determine its dependencies. +transformWithAccInput :: + AliasTable -> + WithAccInput (Aliases GPU) -> + PassM (WithAccInput GPU, Dependencies) +transformWithAccInput aliases (shape, arrs, op) = do + (op', deps) <- case op of + Nothing -> pure (Nothing, mempty) + Just (f, nes) -> do + -- The lambda parameters have no aliases. + (f', deps) <- transformLambda aliases f + pure (Just (f', nes), deps <> depsOf nes) + let deps' = deps <> depsOf shape <> depsOf arrs + pure ((shape, arrs, op'), deps') + +-------------------------------------------------------------------------------- +-- REORDERING - TYPES -- +-------------------------------------------------------------------------------- + +-- | The monad used to reorder statements within a sequence such that its +-- GPUBody statements can be merged into as few possible kernels. +type ReorderM = StateT State PassM + +-- | The state used by a 'ReorderM' monad. +data State = State + { -- | All statements that already have been processed from the sequence, + -- divided into alternating groups of non-GPUBody and GPUBody statements. + -- Groups at even indices only contain non-GPUBody statements. Groups at + -- odd indices only contain GPUBody statements. + stateGroups :: Groups, + stateEquivalents :: EquivalenceTable + } + +-- | A map from variable tags to t'SubExp's returned from within GPUBodies. +type EquivalenceTable = IM.IntMap Entry + +-- | An entry in an 'EquivalenceTable'. +data Entry = Entry + { -- | A value returned from within a GPUBody kernel. + -- In @let res = gpu { x }@ this is @x@. + entryValue :: SubExp, + -- | The type of the 'entryValue'. + entryType :: Type, + -- | The name of the variable that binds the return value for 'entryValue'. + -- In @let res = gpu { x }@ this is @res@. + entryResult :: VName, + -- | The index of the group that `entryResult` is bound in. + entryGroupIdx :: Int, + -- | If 'False' then the entry key is a variable that binds the same value + -- as the 'entryValue'. Otherwise it binds an array with an outer dimension + -- of one whose row equals that value. + entryStored :: Bool + } + +type Groups = SQ.Seq Group + +-- | A group is a subsequence of statements, usually either only GPUBody +-- statements or only non-GPUBody statements. The 'Usage' statistics of those +-- statements are also stored. +data Group = Group + { -- | The statements of the group. + groupStms :: Stms GPU, + -- | The usage statistics of the statements within the group. + groupUsage :: Usage + } + +-- | Usage statistics for some set of statements. +data Usage = Usage + { -- | The variables that the statements bind. + usageBindings :: Bindings, + -- | The variables that the statements depend upon, i.e. the free variables + -- of each statement and the root aliases of every array that they observe. + usageDependencies :: Dependencies + } + +instance Semigroup Group where + (Group s1 u1) <> (Group s2 u2) = Group (s1 <> s2) (u1 <> u2) + +instance Monoid Group where + mempty = Group {groupStms = mempty, groupUsage = mempty} + +instance Semigroup Usage where + (Usage b1 d1) <> (Usage b2 d2) = Usage (b1 <> b2) (d1 <> d2) + +instance Monoid Usage where + mempty = Usage {usageBindings = mempty, usageDependencies = mempty} + +-------------------------------------------------------------------------------- +-- REORDERING - FUNCTIONS -- +-------------------------------------------------------------------------------- + +-- | Return the usage bindings of the group. +groupBindings :: Group -> Bindings +groupBindings = usageBindings . groupUsage + +-- | Return the usage dependencies of the group. +groupDependencies :: Group -> Dependencies +groupDependencies = usageDependencies . groupUsage + +-- | An initial state to use when running a 'ReorderM' monad. +initialState :: State +initialState = + State + { stateGroups = SQ.singleton mempty, + stateEquivalents = mempty + } + +-- | Modify the groups that the sequence has been split into so far. +modifyGroups :: (Groups -> Groups) -> ReorderM () +modifyGroups f = + modify $ \st -> st {stateGroups = f (stateGroups st)} + +-- | Remove these keys from the equivalence table. +removeEquivalents :: IS.IntSet -> ReorderM () +removeEquivalents keys = + modify $ \st -> + let eqs' = stateEquivalents st `IM.withoutKeys` keys + in st {stateEquivalents = eqs'} + +-- | Add an entry to the equivalence table. +recordEquivalent :: VName -> Entry -> ReorderM () +recordEquivalent n entry = + modify $ \st -> + let eqs = stateEquivalents st + eqs' = IM.insert (baseTag n) entry eqs + in st {stateEquivalents = eqs'} + +-- | Moves a GPUBody statement to the furthest possible group of the statement +-- sequence, possibly a new group at the end of sequence. +-- +-- To simplify consumption handling a GPUBody is not allowed to merge with a +-- kernel whose result it consumes. Such GPUBody may therefore not be moved +-- into the same group as such kernel. +moveGPUBody :: Stm GPU -> Usage -> Consumption -> ReorderM () +moveGPUBody stm usage consumed = do + -- Replace dependencies with their GPUBody result equivalents. + eqs <- gets stateEquivalents + let g i = maybe i (baseTag . entryResult) (IM.lookup i eqs) + let deps' = IS.map g (usageDependencies usage) + let usage' = usage {usageDependencies = deps'} + + -- Move the GPUBody. + grps <- gets stateGroups + let f = groupBlocks usage' consumed + let idx = fromMaybe 1 (SQ.findIndexR f grps) + let idx' = case idx `mod` 2 of + 0 -> idx + 1 + _ | consumes idx grps -> idx + 2 + _ -> idx + modifyGroups $ moveToGrp (stm, usage) idx' + + -- Record the kernel equivalents of the bound results. + let pes = patElems (stmPat stm) + let Op (GPUBody _ (Body _ _ res)) = stmExp stm + mapM_ (stores idx') (zip pes (map resSubExp res)) + where + consumes idx grps + | Just grp <- SQ.lookup idx grps = + not $ IS.disjoint (groupBindings grp) consumed + | otherwise = + False + + stores idx (PatElem n t, se) + | Just row_t <- peelArray 1 t = + recordEquivalent n $ Entry se row_t n idx True + | otherwise = + recordEquivalent n $ Entry se t n idx False + +-- | Moves a non-GPUBody statement to the furthest possible groups of the +-- statement sequence, possibly a new group at the end of sequence. +moveOther :: Stm GPU -> Usage -> Consumption -> ReorderM () +moveOther stm usage consumed = do + grps <- gets stateGroups + let f = groupBlocks usage consumed + let idx = fromMaybe 0 (SQ.findIndexR f grps) + let idx' = ((idx + 1) `div` 2) * 2 + modifyGroups $ moveToGrp (stm, usage) idx' + recordEquivalentsOf stm idx' + +-- | @recordEquivalentsOf stm idx@ records the GPUBody result and/or return +-- value that @stm@ is equivalent to. @idx@ is the index of the group that @stm@ +-- belongs to. +-- +-- A GPUBody can have a dependency substituted with a result equivalent if it +-- merges with the source GPUBody, allowing it to be moved beyond the binding +-- site of that dependency. +-- +-- To guarantee that a GPUBody which moves beyond a dependency also merges with +-- its source GPUBody, equivalents are only allowed to be recorded for results +-- bound within the group at index @idx-1@. +recordEquivalentsOf :: Stm GPU -> Int -> ReorderM () +recordEquivalentsOf stm idx = do + eqs <- gets stateEquivalents + case stm of + Let (Pat [PatElem x _]) _ (BasicOp (SubExp (Var n))) + | Just entry <- IM.lookup (baseTag n) eqs, + entryGroupIdx entry == idx - 1 -> + recordEquivalent x entry + Let (Pat [PatElem x _]) _ (BasicOp (Index arr slice)) + | Just entry <- IM.lookup (baseTag arr) eqs, + entryGroupIdx entry == idx - 1, + Slice (DimFix i : dims) <- slice, + i == intConst Int64 0, + dims == map sliceDim (arrayDims $ entryType entry) -> + recordEquivalent x (entry {entryStored = False}) + _ -> pure () + +-- | Does this group block a statement with this usage/consumption statistics +-- from being moved past it? +groupBlocks :: Usage -> Consumption -> Group -> Bool +groupBlocks usage consumed grp = + let bound = groupBindings grp + deps = groupDependencies grp + + used = usageDependencies usage + in not (IS.disjoint bound used && IS.disjoint deps consumed) + +-- | @moveToGrp stm idx grps@ moves @stm@ into the group at index @idx@ of +-- @grps@. +moveToGrp :: (Stm GPU, Usage) -> Int -> Groups -> Groups +moveToGrp stm idx grps + | idx >= SQ.length grps = + moveToGrp stm idx (grps |> mempty) + | otherwise = + SQ.adjust' (stm `moveTo`) idx grps + +-- | Adds the statement and its usage statistics to the group. +moveTo :: (Stm GPU, Usage) -> Group -> Group +moveTo (stm, usage) grp = + grp + { groupStms = groupStms grp |> stm, + groupUsage = groupUsage grp <> usage + } + +-------------------------------------------------------------------------------- +-- MERGING GPU BODIES - TYPES -- +-------------------------------------------------------------------------------- + +-- | The monad used for rewriting a GPUBody to use the t'SubExp's that are +-- returned from kernels it is merged with rather than the results that they +-- bind. +-- +-- The state is a prologue of statements to be added at the beginning of the +-- rewritten kernel body. +type RewriteM = StateT (Stms GPU) ReorderM + +-------------------------------------------------------------------------------- +-- MERGING GPU BODIES - FUNCTIONS -- +-------------------------------------------------------------------------------- + +-- | Collapses the processed sequence of groups into a single group and returns +-- it, merging GPUBody groups into single kernels in the process. +collapse :: ReorderM Group +collapse = do + grps <- zip (cycle [False, True]) . toList <$> gets stateGroups + grp <- foldM clps mempty grps + + modify $ \st -> st {stateGroups = SQ.singleton grp} + pure grp + where + clps grp0 (gpu_bodies, Group stms usage) = do + grp1 <- + if gpu_bodies + then Group <$> mergeKernels stms <*> pure usage + else pure (Group stms usage) + -- Remove equivalents that no longer are relevant for rewriting GPUBody + -- kernels. This ensures that they are not substituted in later kernels + -- where the replacement variables might not be in scope. + removeEquivalents (groupBindings grp1) + pure (grp0 <> grp1) + +-- | Merges a sequence of GPUBody statements into a single kernel. +mergeKernels :: Stms GPU -> ReorderM (Stms GPU) +mergeKernels stms + | SQ.length stms < 2 = + pure stms + | otherwise = + SQ.singleton <$> foldrM merge empty stms + where + empty = Let mempty (StmAux mempty mempty ()) noop + noop = Op (GPUBody [] (Body () SQ.empty [])) + + merge :: Stm GPU -> Stm GPU -> ReorderM (Stm GPU) + merge stm0 stm1 + | Let pat0 (StmAux cs0 attrs0 _) (Op (GPUBody types0 body)) <- stm0, + Let pat1 (StmAux cs1 attrs1 _) (Op (GPUBody types1 body1)) <- stm1 = + do + Body _ stms0 res0 <- execRewrite (rewriteBody body) + let Body _ stms1 res1 = body1 + + pat' = pat0 <> pat1 + aux' = StmAux (cs0 <> cs1) (attrs0 <> attrs1) () + types' = types0 ++ types1 + body' = Body () (stms0 <> stms1) (res0 <> res1) + in pure (Let pat' aux' (Op (GPUBody types' body'))) + merge _ _ = + compilerBugS "mergeGPUBodies: cannot merge non-GPUBody statements" + +-- | Perform a rewrite and finish it by adding the rewrite prologue to the start +-- of the body. +execRewrite :: RewriteM (Body GPU) -> ReorderM (Body GPU) +execRewrite m = evalStateT m' SQ.empty + where + m' = do + Body _ stms res <- m + prologue <- get + pure (Body () (prologue <> stms) res) + +-- | Return the equivalence table. +equivalents :: RewriteM EquivalenceTable +equivalents = lift (gets stateEquivalents) + +rewriteBody :: Body GPU -> RewriteM (Body GPU) +rewriteBody (Body _ stms res) = + Body () <$> rewriteStms stms <*> rewriteResult res + +rewriteStms :: Stms GPU -> RewriteM (Stms GPU) +rewriteStms = mapM rewriteStm + +rewriteStm :: Stm GPU -> RewriteM (Stm GPU) +rewriteStm (Let (Pat pes) (StmAux cs attrs _) e) = do + pat' <- Pat <$> mapM rewritePatElem pes + cs' <- rewriteCerts cs + e' <- rewriteExp e + pure $ Let pat' (StmAux cs' attrs ()) e' + +rewritePatElem :: PatElem Type -> RewriteM (PatElem Type) +rewritePatElem (PatElem n t) = + PatElem n <$> rewriteType t + +rewriteExp :: Exp GPU -> RewriteM (Exp GPU) +rewriteExp e = do + eqs <- equivalents + case e of + BasicOp (Index arr slice) + | Just entry <- IM.lookup (baseTag arr) eqs, + DimFix idx : dims <- unSlice slice, + idx == intConst Int64 0 -> + let se = entryValue entry + in pure . BasicOp $ case (dims, se) of + ([], _) -> SubExp se + (_, Var src) -> Index src (Slice dims) + _ -> compilerBugS "rewriteExp: bad equivalence entry" + _ -> mapExpM rewriter e + where + rewriter = + Mapper + { mapOnSubExp = rewriteSubExp, + mapOnBody = const rewriteBody, + mapOnVName = rewriteName, + mapOnRetType = rewriteExtType, + mapOnBranchType = rewriteExtType, + mapOnFParam = rewriteParam, + mapOnLParam = rewriteParam, + mapOnOp = const opError + } + + opError = compilerBugS "rewriteExp: unhandled HostOp in GPUBody" + +rewriteResult :: Result -> RewriteM Result +rewriteResult = mapM rewriteSubExpRes + +rewriteSubExpRes :: SubExpRes -> RewriteM SubExpRes +rewriteSubExpRes (SubExpRes cs se) = + SubExpRes <$> rewriteCerts cs <*> rewriteSubExp se + +rewriteCerts :: Certs -> RewriteM Certs +rewriteCerts (Certs cs) = + Certs <$> mapM rewriteName cs + +rewriteType :: TypeBase Shape u -> RewriteM (TypeBase Shape u) +-- Note: mapOnType also maps the VName token of accumulators +rewriteType = mapOnType rewriteSubExp + +rewriteExtType :: TypeBase ExtShape u -> RewriteM (TypeBase ExtShape u) +-- Note: mapOnExtType also maps the VName token of accumulators +rewriteExtType = mapOnExtType rewriteSubExp + +rewriteParam :: Param (TypeBase Shape u) -> RewriteM (Param (TypeBase Shape u)) +rewriteParam (Param attrs n t) = + Param attrs n <$> rewriteType t + +rewriteSubExp :: SubExp -> RewriteM SubExp +rewriteSubExp (Constant c) = pure (Constant c) +rewriteSubExp (Var n) = do + eqs <- equivalents + case IM.lookup (baseTag n) eqs of + Nothing -> pure (Var n) + Just (Entry se _ _ _ False) -> pure se + Just (Entry se t _ _ True) -> Var <$> asArray se t + +rewriteName :: VName -> RewriteM VName +rewriteName n = do + se <- rewriteSubExp (Var n) + case se of + Var n' -> pure n' + Constant c -> referConst c + +-- | @asArray se t@ adds @let x = [se]@ to the rewrite prologue and returns the +-- name of @x@. @t@ is the type of @se@. +asArray :: SubExp -> Type -> RewriteM VName +asArray se row_t = do + name <- newName "arr" + let t = row_t `arrayOfRow` intConst Int64 1 + + let pat = Pat [PatElem name t] + let aux = StmAux mempty mempty () + let e = BasicOp (ArrayLit [se] row_t) + + modify (|> Let pat aux e) + pure name + +-- | @referConst c@ adds @let x = c@ to the rewrite prologue and returns the +-- name of @x@. +referConst :: PrimValue -> RewriteM VName +referConst c = do + name <- newName "cnst" + let t = Prim (primValueType c) + + let pat = Pat [PatElem name t] + let aux = StmAux mempty mempty () + let e = BasicOp (SubExp $ Constant c) + + modify (|> Let pat aux e) + pure name + +-- | Produce a fresh name, using the given string as a template. +newName :: String -> RewriteM VName +newName s = lift $ lift (newNameFromString s) diff --git a/src/Futhark/Optimise/ReduceDeviceSyncs.hs b/src/Futhark/Optimise/ReduceDeviceSyncs.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/ReduceDeviceSyncs.hs @@ -0,0 +1,848 @@ +-- | This module implements an optimization that migrates host +-- statements into 'GPUBody' kernels to reduce the number of +-- host-device synchronizations that occur when a scalar variable is +-- written to or read from device memory. Which statements that should +-- be migrated are determined by a 'MigrationTable' produced by the +-- "Futhark.Optimise.ReduceDeviceSyncs.MigrationTable" module; this module +-- merely performs the migration and rewriting dictated by that table. +module Futhark.Optimise.ReduceDeviceSyncs (reduceDeviceSyncs) where + +import Control.Monad +import Control.Monad.Trans.Class +import qualified Control.Monad.Trans.Reader as R +import Control.Monad.Trans.State.Strict hiding (State) +import Control.Parallel.Strategies (parMap, rpar) +import Data.Foldable +import qualified Data.IntMap.Strict as IM +import Data.List (unzip4, zip4) +import qualified Data.Map.Strict as M +import Data.Sequence ((<|), (><), (|>)) +import qualified Data.Text as T +import Futhark.Construct (fullSlice, sliceDim) +import Futhark.Error +import qualified Futhark.FreshNames as FN +import Futhark.IR.GPU +import Futhark.MonadFreshNames (VNameSource, getNameSource, putNameSource) +import Futhark.Optimise.ReduceDeviceSyncs.MigrationTable +import Futhark.Pass +import Futhark.Transform.Substitute + +-- | An optimization pass that migrates host statements into 'GPUBody' kernels +-- to reduce the number of host-device synchronizations. +reduceDeviceSyncs :: Pass GPU GPU +reduceDeviceSyncs = + Pass + "reduce device synchronizations" + "Move host statements to device to reduce blocking memory operations." + run + where + run prog = do + ns <- getNameSource + let mt = analyseProg prog + let st = initialState ns + let (prog', st') = R.runReader (runStateT (optimizeProgram prog) st) mt + putNameSource (stateNameSource st') + pure prog' + +-------------------------------------------------------------------------------- +-- AD HOC OPTIMIZATION -- +-------------------------------------------------------------------------------- + +-- | Optimize a whole program. The type signatures of top-level functions will +-- remain unchanged. +optimizeProgram :: Prog GPU -> ReduceM (Prog GPU) +optimizeProgram (Prog consts funs) = do + consts' <- optimizeStms consts + funs' <- sequence $ parMap rpar optimizeFunDef funs + pure (Prog consts' funs') + +-- | Optimize a function definition. Its type signature will remain unchanged. +optimizeFunDef :: FunDef GPU -> ReduceM (FunDef GPU) +optimizeFunDef fd = do + let body = funDefBody fd + stms' <- optimizeStms (bodyStms body) + pure $ fd {funDefBody = body {bodyStms = stms'}} + +-- | Optimize a body. Scalar results may be replaced with single-element arrays. +optimizeBody :: Body GPU -> ReduceM (Body GPU) +optimizeBody (Body _ stms res) = do + stms' <- optimizeStms stms + res' <- resolveResult res + pure (Body () stms' res') + +-- | Optimize a sequence of statements. +optimizeStms :: Stms GPU -> ReduceM (Stms GPU) +optimizeStms = foldM optimizeStm mempty + +-- | Optimize a single statement, rewriting it into one or more statements to +-- be appended to the provided 'Stms'. Only variables with continued host usage +-- will remain in scope if their statement is migrated. +optimizeStm :: Stms GPU -> Stm GPU -> ReduceM (Stms GPU) +optimizeStm out stm = do + move <- asks (shouldMoveStm stm) + if move + then moveStm out stm + else case stmExp stm of + BasicOp (Update safety arr slice (Var v)) + | Just _ <- sliceIndices slice -> do + -- Rewrite the Update if its write value has been migrated. Copying + -- is faster than doing a synchronous write, so we use the device + -- array even if the value has been made available to the host. + dev <- storedScalar (Var v) + case dev of + Nothing -> pure (out |> stm) + Just dst -> do + -- Transform the single element Update into a slice Update. + let dims = unSlice slice + let (outer, [DimFix i]) = splitAt (length dims - 1) dims + let one = intConst Int64 1 + let slice' = Slice $ outer ++ [DimSlice i one one] + let e = BasicOp (Update safety arr slice' (Var dst)) + let stm' = stm {stmExp = e} + + pure (out |> stm') + BasicOp (Replicate (Shape dims) (Var v)) + | Pat [PatElem n arr_t] <- stmPat stm -> do + -- A Replicate can be rewritten to not require its replication value + -- to be available on host. If its value is migrated the Replicate + -- thus needs to be transformed. + -- + -- If the inner dimension of the replication array is one then the + -- rewrite can be performed more efficiently than the general case. + v' <- resolveName v + let v_kept_on_device = v /= v' + + gpubody_ok <- gets stateGPUBodyOk + + case v_kept_on_device of + False -> pure (out |> stm) + True + | all (== intConst Int64 1) dims, + Just t' <- peelArray 1 arr_t, + gpubody_ok -> do + let n' = VName (baseName n `withSuffix` "_inner") 0 + let pat' = Pat [PatElem n' t'] + let e' = BasicOp $ Replicate (Shape $ tail dims) (Var v) + let stm' = Let pat' (stmAux stm) e' + + -- `gpu { v }` is slightly faster than `replicate 1 v` and + -- can merge with the GPUBody that v was computed by. + gpubody <- inGPUBody (rewriteStm stm') + pure (out |> gpubody {stmPat = stmPat stm}) + True + | last dims == intConst Int64 1 -> + let e' = BasicOp $ Replicate (Shape $ init dims) (Var v') + stm' = stm {stmExp = e'} + in pure (out |> stm') + True -> do + n' <- newName n + -- v_kept_on_device implies that v is a scalar. + let dims' = dims ++ [intConst Int64 1] + let arr_t' = Array (elemType arr_t) (Shape dims') NoUniqueness + let pat' = Pat [PatElem n' arr_t'] + let e' = BasicOp $ Replicate (Shape dims) (Var v') + let repl = Let pat' (stmAux stm) e' + + let aux = StmAux mempty mempty () + let slice = map sliceDim (arrayDims arr_t) + let slice' = slice ++ [DimFix $ intConst Int64 0] + let idx = BasicOp $ Index n' (Slice slice') + let index = Let (stmPat stm) aux idx + + pure (out |> repl |> index) + BasicOp {} -> + pure (out |> stm) + Apply {} -> + pure (out |> stm) + If cond (Body _ tstms0 tres) (Body _ fstms0 fres) (IfDec btypes sort) -> + do + -- Rewrite branches. + tstms1 <- optimizeStms tstms0 + fstms1 <- optimizeStms fstms0 + + -- Ensure return values and types match if one or both branches + -- return a result that now reside on device. + let bmerge (res, tstms, fstms) (pe, tr, fr, bt) = + do + let onHost (Var v) = (v ==) <$> resolveName v + onHost _ = pure True + + tr_on_host <- onHost (resSubExp tr) + fr_on_host <- onHost (resSubExp fr) + + if tr_on_host && fr_on_host + then -- No result resides on device ==> nothing to do. + pure ((pe, tr, fr, bt) : res, tstms, fstms) + else -- Otherwise, ensure both results are migrated. + do + let t = patElemType pe + (tstms', tarr) <- storeScalar tstms (resSubExp tr) t + (fstms', farr) <- storeScalar fstms (resSubExp fr) t + + pe' <- arrayizePatElem pe + let bt' = staticShapes1 (patElemType pe') + let tr' = tr {resSubExp = Var tarr} + let fr' = fr {resSubExp = Var farr} + pure ((pe', tr', fr', bt') : res, tstms', fstms') + + let pes = patElems (stmPat stm) + let zipped = zip4 pes tres fres btypes + (zipped', tstms2, fstms2) <- foldM bmerge ([], tstms1, fstms1) zipped + let (pes', tres', fres', btypes') = unzip4 (reverse zipped') + + -- Rewrite statement. + let tbranch' = Body () tstms2 tres' + let fbranch' = Body () fstms2 fres' + let e' = If cond tbranch' fbranch' (IfDec btypes' sort) + let stm' = Let (Pat pes') (stmAux stm) e' + + -- Read migrated scalars that are used on host. + foldM addRead (out |> stm') (zip pes pes') + DoLoop ps lf b -> do + -- Enable the migration of for-in loop variables. + (params, lform, body) <- rewriteForIn (ps, lf, b) + + -- Update statement bound variables and parameters if their values + -- have been migrated to device. + let lmerge (res, stms) (pe, (Param _ pn pt, pval), MoveToDevice) = do + -- Rewrite the bound variable. + pe' <- arrayizePatElem pe + + -- Move the initial value to device if not already there. + (stms', arr) <- storeScalar stms pval (fromDecl pt) + + -- Rewrite the parameter. + pn' <- newName pn + let pt' = toDecl (patElemType pe') Nonunique + let pval' = Var arr + let param' = (Param mempty pn' pt', pval') + + -- Record the migration. + Ident pn (fromDecl pt) `movedTo` pn' + + pure ((pe', param') : res, stms') + lmerge _ (_, _, UsedOnHost) = + -- Initial loop parameter value and loop result should have + -- been made available on host instead. + compilerBugS "optimizeStm: unhandled host usage of loop param" + lmerge (res, stms) (pe, param, StayOnHost) = + pure ((pe, param) : res, stms) + + mt <- ask + + let pes = patElems (stmPat stm) + let mss = map (\(Param _ n _, _) -> statusOf n mt) params + (zipped', out') <- foldM lmerge ([], out) (zip3 pes params mss) + let (pes', params') = unzip (reverse zipped') + + -- Rewrite statement. + body' <- optimizeBody body + let e' = DoLoop params' lform body' + let stm' = Let (Pat pes') (stmAux stm) e' + + -- Read migrated scalars that are used on host. + foldM addRead (out' |> stm') (zip pes pes') + WithAcc inputs lmd -> do + let getAcc (Acc a _ _ _) = a + getAcc _ = + compilerBugS + "Type error: WithAcc expression did not return accumulator." + + let accs = zipWith (\t i -> (getAcc t, i)) (lambdaReturnType lmd) inputs + inputs' <- mapM (uncurry optimizeWithAccInput) accs + + let body = lambdaBody lmd + stms' <- optimizeStms (bodyStms body) + + let rewrite (SubExpRes certs se, t, pe) = + do + se' <- resolveSubExp se + if se == se' + then pure (SubExpRes certs se, t, pe) + else do + pe' <- arrayizePatElem pe + let t' = patElemType pe' + pure (SubExpRes certs se', t', pe') + + -- Rewrite non-accumulator results that have been migrated. + -- + -- Accumulator return values do not map to arrays one-to-one but + -- one-to-many. They are not transformed however and can be mapped + -- as a no-op. + let len = length inputs + let (res0, res1) = splitAt len (bodyResult body) + let (rts0, rts1) = splitAt len (lambdaReturnType lmd) + let pes = patElems (stmPat stm) + let (pes0, pes1) = splitAt (length pes - length res1) pes + (res1', rts1', pes1') <- unzip3 <$> mapM rewrite (zip3 res1 rts1 pes1) + let res' = res0 ++ res1' + let rts' = rts0 ++ rts1' + let pes' = pes0 ++ pes1' + + -- Rewrite statement. + let body' = Body () stms' res' + let lmd' = lmd {lambdaBody = body', lambdaReturnType = rts'} + let e' = WithAcc inputs' lmd' + let stm' = Let (Pat pes') (stmAux stm) e' + + -- Read migrated scalars that are used on host. + foldM addRead (out |> stm') (zip pes pes') + Op op -> do + op' <- optimizeHostOp op + pure (out |> stm {stmExp = Op op'}) + where + addRead stms (pe@(PatElem n _), PatElem dev _) + | n == dev = pure stms + | otherwise = pe `migratedTo` (dev, stms) + +-- | Rewrite a for-in loop such that relevant source array reads can be delayed. +rewriteForIn :: + ([(FParam GPU, SubExp)], LoopForm GPU, Body GPU) -> + ReduceM ([(FParam GPU, SubExp)], LoopForm GPU, Body GPU) +rewriteForIn loop@(_, WhileLoop {}, _) = + pure loop +rewriteForIn (params, ForLoop i t n elems, body) = do + mt <- ask + let (elems', stms') = foldr (inline mt) ([], bodyStms body) elems + pure (params, ForLoop i t n elems', body {bodyStms = stms'}) + where + inline mt (x, arr) (arrs, stms) + | pn <- paramName x, + not (usedOnHost pn mt) = + let pt = typeOf x + stm = bind (PatElem pn pt) (BasicOp $ index arr pt) + in (arrs, stm <| stms) + | otherwise = + ((x, arr) : arrs, stms) + + index arr of_type = + Index arr $ Slice $ DimFix (Var i) : map sliceDim (arrayDims of_type) + +-- | Optimize an accumulator input. The 'VName' is the accumulator token. +optimizeWithAccInput :: VName -> WithAccInput GPU -> ReduceM (WithAccInput GPU) +optimizeWithAccInput _ (shape, arrs, Nothing) = pure (shape, arrs, Nothing) +optimizeWithAccInput acc (shape, arrs, Just (op, nes)) = do + device_only <- asks (shouldMove acc) + if device_only + then do + op' <- addReadsToLambda op + pure (shape, arrs, Just (op', nes)) + else do + let body = lambdaBody op + -- To pass type check neither parameters nor results can change. + -- + -- op may be used on both host and device so we must avoid introducing + -- any GPUBody statements. + stms' <- noGPUBody $ optimizeStms (bodyStms body) + let op' = op {lambdaBody = body {bodyStms = stms'}} + pure (shape, arrs, Just (op', nes)) + +-- | Optimize a host operation. 'Index' statements are added to kernel code +-- that depends on migrated scalars. +optimizeHostOp :: HostOp GPU op -> ReduceM (HostOp GPU op) +optimizeHostOp (SegOp (SegMap lvl space types kbody)) = + SegOp . SegMap lvl space types <$> addReadsToKernelBody kbody +optimizeHostOp (SegOp (SegRed lvl space ops types kbody)) = do + ops' <- mapM addReadsToSegBinOp ops + SegOp . SegRed lvl space ops' types <$> addReadsToKernelBody kbody +optimizeHostOp (SegOp (SegScan lvl space ops types kbody)) = do + ops' <- mapM addReadsToSegBinOp ops + SegOp . SegScan lvl space ops' types <$> addReadsToKernelBody kbody +optimizeHostOp (SegOp (SegHist lvl space ops types kbody)) = do + ops' <- mapM addReadsToHistOp ops + SegOp . SegHist lvl space ops' types <$> addReadsToKernelBody kbody +optimizeHostOp (SizeOp op) = + pure (SizeOp op) +optimizeHostOp OtherOp {} = + -- These should all have been taken care of in the unstreamGPU pass. + compilerBugS "optimizeHostOp: unhandled OtherOp" +optimizeHostOp (GPUBody types body) = + GPUBody types <$> addReadsToBody body + +-------------------------------------------------------------------------------- +-- COMMON HELPERS -- +-------------------------------------------------------------------------------- + +-- | Append the given string to a name. +withSuffix :: Name -> String -> Name +withSuffix name sfx = nameFromText $ T.append (nameToText name) (T.pack sfx) + +-------------------------------------------------------------------------------- +-- MIGRATION - TYPES -- +-------------------------------------------------------------------------------- + +-- | The monad used to perform migration-based synchronization reductions. +type ReduceM = StateT State (R.Reader MigrationTable) + +-- | The state used by a 'ReduceM' monad. +data State = State + { -- | A source to generate new 'VName's from. + stateNameSource :: VNameSource, + -- | A table of variables in the original program which have been migrated + -- to device. Each variable maps to a tuple that describes: + -- * 'baseName' of the original variable. + -- * Type of the original variable. + -- * Name of the single element array holding the migrated value. + -- * Whether the original variable still can be used on the host. + stateMigrated :: IM.IntMap (Name, Type, VName, Bool), + -- | Whether non-migration optimizations may introduce 'GPUBody' kernels at + -- the current location. + stateGPUBodyOk :: Bool + } + +-------------------------------------------------------------------------------- +-- MIGRATION - PRIMITIVES -- +-------------------------------------------------------------------------------- + +-- | An initial state to use when running a 'ReduceM' monad. +initialState :: VNameSource -> State +initialState ns = + State + { stateNameSource = ns, + stateMigrated = mempty, + stateGPUBodyOk = True + } + +-- | Retrieve a function of the current environment. +asks :: (MigrationTable -> a) -> ReduceM a +asks = lift . R.asks + +-- | Fetch the value of the environment. +ask :: ReduceM MigrationTable +ask = lift R.ask + +-- | Perform non-migration optimizations without introducing any GPUBody +-- kernels. +noGPUBody :: ReduceM a -> ReduceM a +noGPUBody m = do + prev <- gets stateGPUBodyOk + modify $ \st -> st {stateGPUBodyOk = False} + res <- m + modify $ \st -> st {stateGPUBodyOk = prev} + pure res + +-- | Produce a fresh name, using the given name as a template. +newName :: VName -> ReduceM VName +newName n = do + st <- get + let ns = stateNameSource st + let (n', ns') = FN.newName ns n + put (st {stateNameSource = ns'}) + pure n' + +-- | Create a 'PatElem' that binds the array of a migrated variable binding. +arrayizePatElem :: PatElem Type -> ReduceM (PatElem Type) +arrayizePatElem (PatElem n t) = do + let name = baseName n `withSuffix` "_dev" + dev <- newName (VName name 0) + let dev_t = t `arrayOfRow` intConst Int64 1 + pure (PatElem dev dev_t) + +-- | @x `movedTo` arr@ registers that the value of @x@ has been migrated to +-- @arr[0]@. +movedTo :: Ident -> VName -> ReduceM () +movedTo = recordMigration False + +-- | @x `aliasedBy` arr@ registers that the value of @x@ also is available on +-- device as @arr[0]@. +aliasedBy :: Ident -> VName -> ReduceM () +aliasedBy = recordMigration True + +-- | @recordMigration host x arr@ records the migration of variable @x@ to +-- @arr[0]@. If @host@ then the original binding can still be used on host. +recordMigration :: Bool -> Ident -> VName -> ReduceM () +recordMigration host (Ident x t) arr = + modify $ \st -> + let migrated = stateMigrated st + entry = (baseName x, t, arr, host) + migrated' = IM.insert (baseTag x) entry migrated + in st {stateMigrated = migrated'} + +-- | @pe `migratedTo` (dev, stms)@ registers that the variable @pe@ in the +-- original program has been migrated to @dev@ and rebinds the variable if +-- deemed necessary, adding an index statement to the given statements. +migratedTo :: PatElem Type -> (VName, Stms GPU) -> ReduceM (Stms GPU) +migratedTo pe (dev, stms) = do + used <- asks (usedOnHost $ patElemName pe) + if used + then patElemIdent pe `aliasedBy` dev >> pure (stms |> bind pe (eIndex dev)) + else patElemIdent pe `movedTo` dev >> pure stms + +-- | @useScalar stms n@ returns a variable that binds the result bound by @n@ +-- in the original program. If the variable has been migrated to device and have +-- not been copied back to host a new variable binding will be added to the +-- provided statements and be returned. +useScalar :: Stms GPU -> VName -> ReduceM (Stms GPU, VName) +useScalar stms n = do + entry <- IM.lookup (baseTag n) <$> gets stateMigrated + case entry of + Nothing -> + pure (stms, n) + Just (_, _, _, True) -> + pure (stms, n) + Just (name, t, arr, _) -> + do + n' <- newName (VName name 0) + let stm = bind (PatElem n' t) (eIndex arr) + pure (stms |> stm, n') + +-- | Create an expression that reads the first element of a 1-dimensional array. +eIndex :: VName -> Exp GPU +eIndex arr = BasicOp $ Index arr (Slice [DimFix $ intConst Int64 0]) + +-- | A shorthand for binding a single variable to an expression. +bind :: PatElem Type -> Exp GPU -> Stm GPU +bind pe = Let (Pat [pe]) (StmAux mempty mempty ()) + +-- | Returns the array alias of @se@ if it is a variable that has been migrated +-- to device. Otherwise returns @Nothing@. +storedScalar :: SubExp -> ReduceM (Maybe VName) +storedScalar (Constant _) = pure Nothing +storedScalar (Var n) = do + entry <- IM.lookup (baseTag n) <$> gets stateMigrated + pure $ fmap (\(_, _, arr, _) -> arr) entry + +-- | @storeScalar stms se t@ returns a variable that binds a single element +-- array that contains the value of @se@ in the original program. If @se@ is a +-- variable that has been migrated to device, its existing array alias will be +-- used. Otherwise a new variable binding will be added to the provided +-- statements and be returned. @t@ is the type of @se@. +storeScalar :: Stms GPU -> SubExp -> Type -> ReduceM (Stms GPU, VName) +storeScalar stms se t = do + entry <- case se of + Var n -> IM.lookup (baseTag n) <$> gets stateMigrated + _ -> pure Nothing + case entry of + Just (_, _, arr, _) -> pure (stms, arr) + Nothing -> do + -- How to most efficiently create an array containing the given value + -- depends on whether it is a variable or a constant. Creating a constant + -- array is a runtime copy of static memory, while creating an array that + -- contains a variable results in a synchronous write. The latter is thus + -- replaced with either a mergeable GPUBody kernel or a Replicate. + -- + -- Whether it makes sense to hoist arrays out of bodies to enable CSE is + -- left to the simplifier to figure out. Duplicates will be eliminated if + -- a scalar is stored multiple times within a body. + -- + -- TODO: Enable the simplifier to hoist non-consumed, non-returned arrays + -- out of top-level function definitions. All constant arrays + -- produced here are in principle top-level hoistable. + gpubody_ok <- gets stateGPUBodyOk + case se of + Var n | gpubody_ok -> do + n' <- newName n + let stm = bind (PatElem n' t) (BasicOp $ SubExp se) + + gpubody <- inGPUBody (pure stm) + let dev = patElemName $ head $ patElems (stmPat gpubody) + + pure (stms |> gpubody, dev) + Var n -> do + pe <- arrayizePatElem (PatElem n t) + let shape = Shape [intConst Int64 1] + let stm = bind pe (BasicOp $ Replicate shape se) + pure (stms |> stm, patElemName pe) + _ -> do + let n = VName (nameFromString "const") 0 + pe <- arrayizePatElem (PatElem n t) + let stm = bind pe (BasicOp $ ArrayLit [se] t) + pure (stms |> stm, patElemName pe) + +-- | Map a variable name to itself or, if the variable no longer can be used on +-- host, the name of a single element array containing its value. +resolveName :: VName -> ReduceM VName +resolveName n = do + entry <- IM.lookup (baseTag n) <$> gets stateMigrated + case entry of + Nothing -> pure n + Just (_, _, _, True) -> pure n + Just (_, _, arr, _) -> pure arr + +-- | Like 'resolveName' but for a t'SubExp'. Constants are mapped to themselves. +resolveSubExp :: SubExp -> ReduceM SubExp +resolveSubExp (Var n) = Var <$> resolveName n +resolveSubExp cnst = pure cnst + +-- | Like 'resolveSubExp' but for a 'SubExpRes'. +resolveSubExpRes :: SubExpRes -> ReduceM SubExpRes +resolveSubExpRes (SubExpRes certs se) = + -- Certificates are always read back to host. + SubExpRes certs <$> resolveSubExp se + +-- | Apply 'resolveSubExpRes' to a list of results. +resolveResult :: Result -> ReduceM Result +resolveResult = mapM resolveSubExpRes + +-- | Migrate a statement to device, ensuring all its bound variables used on +-- host will remain available with the same names. +moveStm :: Stms GPU -> Stm GPU -> ReduceM (Stms GPU) +moveStm out (Let pat aux (BasicOp (ArrayLit [se] t'))) + | Pat [PatElem n _] <- pat = + do + -- Save an 'Index' by rewriting the 'ArrayLit' rather than migrating it. + let n' = VName (baseName n `withSuffix` "_inner") 0 + let pat' = Pat [PatElem n' t'] + let e' = BasicOp (SubExp se) + let stm' = Let pat' aux e' + + gpubody <- inGPUBody (rewriteStm stm') + pure (out |> gpubody {stmPat = pat}) +moveStm out stm = do + -- Move the statement to device. + gpubody <- inGPUBody (rewriteStm stm) + + -- Read non-scalars and scalars that are used on host. + let arrs = zip (patElems $ stmPat stm) (patElems $ stmPat gpubody) + foldM addRead (out |> gpubody) arrs + where + addRead stms (pe@(PatElem _ t), PatElem dev dev_t) = + let add' e = pure $ stms |> bind pe e + add = add' . BasicOp + in case arrayRank dev_t of + -- Alias non-arrays with their prior name. + 0 -> add $ SubExp (Var dev) + -- Read all certificates for free. + 1 | t == Prim Unit -> add' (eIndex dev) + -- Record the device alias of each scalar variable and read them + -- if used on host. + 1 -> pe `migratedTo` (dev, stms) + -- Drop the added dimension of multidimensional arrays. + _ -> add $ Index dev (fullSlice dev_t [DimFix $ intConst Int64 0]) + +-- | Create a GPUBody kernel that executes a single statement and stores its +-- results in single element arrays. +inGPUBody :: RewriteM (Stm GPU) -> ReduceM (Stm GPU) +inGPUBody m = do + (stm, st) <- runStateT m initialRState + let prologue = rewritePrologue st + + let pes = patElems (stmPat stm) + pat <- Pat <$> mapM arrayizePatElem pes + let aux = StmAux mempty mempty () + let types = map patElemType pes + let res = map (SubExpRes mempty . Var . patElemName) pes + let body = Body () (prologue |> stm) res + let e = Op (GPUBody types body) + pure (Let pat aux e) + +-------------------------------------------------------------------------------- +-- KERNEL REWRITING - TYPES -- +-------------------------------------------------------------------------------- + +-- The monad used to rewrite (migrated) kernel code. +type RewriteM = StateT RState ReduceM + +-- | The state used by a 'RewriteM' monad. +data RState = RState + { -- | Maps variables in the original program to names to be used by rewrites. + rewriteRenames :: IM.IntMap VName, + -- | Statements to be added as a prologue before rewritten statements. + rewritePrologue :: Stms GPU + } + +-------------------------------------------------------------------------------- +-- KERNEL REWRITING - FUNCTIONS -- +-------------------------------------------------------------------------------- + +-- | An initial state to use when running a 'RewriteM' monad. +initialRState :: RState +initialRState = + RState + { rewriteRenames = mempty, + rewritePrologue = mempty + } + +-- | Rewrite 'SegBinOp' dependencies to scalars that have been migrated. +addReadsToSegBinOp :: SegBinOp GPU -> ReduceM (SegBinOp GPU) +addReadsToSegBinOp op = do + f' <- addReadsToLambda (segBinOpLambda op) + pure (op {segBinOpLambda = f'}) + +-- | Rewrite 'HistOp' dependencies to scalars that have been migrated. +addReadsToHistOp :: HistOp GPU -> ReduceM (HistOp GPU) +addReadsToHistOp op = do + f' <- addReadsToLambda (histOp op) + pure (op {histOp = f'}) + +-- | Rewrite generic lambda dependencies to scalars that have been migrated. +addReadsToLambda :: Lambda GPU -> ReduceM (Lambda GPU) +addReadsToLambda f = do + body' <- addReadsToBody (lambdaBody f) + pure (f {lambdaBody = body'}) + +-- | Rewrite generic body dependencies to scalars that have been migrated. +addReadsToBody :: Body GPU -> ReduceM (Body GPU) +addReadsToBody body = do + (body', prologue) <- addReadsHelper body + pure body' {bodyStms = prologue >< bodyStms body'} + +-- | Rewrite kernel body dependencies to scalars that have been migrated. +addReadsToKernelBody :: KernelBody GPU -> ReduceM (KernelBody GPU) +addReadsToKernelBody kbody = do + (kbody', prologue) <- addReadsHelper kbody + pure kbody' {kernelBodyStms = prologue >< kernelBodyStms kbody'} + +-- | Rewrite migrated scalar dependencies within anything. The returned +-- statements must be added to the scope of the rewritten construct. +addReadsHelper :: (FreeIn a, Substitute a) => a -> ReduceM (a, Stms GPU) +addReadsHelper x = do + let from = namesToList (freeIn x) + (to, st) <- runStateT (mapM rename from) initialRState + let rename_map = M.fromList (zip from to) + pure (substituteNames rename_map x, rewritePrologue st) + +-- | Create a fresh name, registering which name it is a rewrite of. +rewriteName :: VName -> RewriteM VName +rewriteName n = do + n' <- lift (newName n) + modify $ \st -> st {rewriteRenames = IM.insert (baseTag n) n' (rewriteRenames st)} + pure n' + +-- | Rewrite all bindings introduced by a body (to ensure they are unique) and +-- fix any dependencies that are broken as a result of migration or rewriting. +rewriteBody :: Body GPU -> RewriteM (Body GPU) +rewriteBody (Body _ stms res) = do + stms' <- rewriteStms stms + res' <- renameResult res + pure (Body () stms' res') + +-- | Rewrite all bindings introduced by a sequence of statements (to ensure they +-- are unique) and fix any dependencies that are broken as a result of migration +-- or rewriting. +rewriteStms :: Stms GPU -> RewriteM (Stms GPU) +rewriteStms = foldM rewriteTo mempty + where + rewriteTo out stm = do + stm' <- rewriteStm stm + pure $ case stmExp stm' of + Op (GPUBody _ (Body _ stms res)) -> + let pes = patElems (stmPat stm') + in foldl' bnd (out >< stms) (zip pes res) + _ -> out |> stm' + + bnd :: Stms GPU -> (PatElem Type, SubExpRes) -> Stms GPU + bnd out (pe, SubExpRes cs se) + | Just t' <- peelArray 1 (typeOf pe) = + out |> Let (Pat [pe]) (StmAux cs mempty ()) (BasicOp $ ArrayLit [se] t') + | otherwise = + out |> Let (Pat [pe]) (StmAux cs mempty ()) (BasicOp $ SubExp se) + +-- | Rewrite all bindings introduced by a single statement (to ensure they are +-- unique) and fix any dependencies that are broken as a result of migration or +-- rewriting. +-- +-- NOTE: GPUBody kernels must be rewritten through 'rewriteStms'. +rewriteStm :: Stm GPU -> RewriteM (Stm GPU) +rewriteStm (Let pat aux e) = do + e' <- rewriteExp e + pat' <- rewritePat pat + aux' <- rewriteStmAux aux + pure (Let pat' aux' e') + +-- | Rewrite all bindings introduced by a pattern (to ensure they are unique) +-- and fix any dependencies that are broken as a result of migration or +-- rewriting. +rewritePat :: Pat Type -> RewriteM (Pat Type) +rewritePat pat = Pat <$> mapM rewritePatElem (patElems pat) + +-- | Rewrite the binding introduced by a single pattern element (to ensure it is +-- unique) and fix any dependencies that are broken as a result of migration or +-- rewriting. +rewritePatElem :: PatElem Type -> RewriteM (PatElem Type) +rewritePatElem (PatElem n t) = do + n' <- rewriteName n + t' <- renameType t + pure (PatElem n' t') + +-- | Fix any 'StmAux' certificate references that are broken as a result of +-- migration or rewriting. +rewriteStmAux :: StmAux () -> RewriteM (StmAux ()) +rewriteStmAux (StmAux certs attrs _) = do + certs' <- renameCerts certs + pure (StmAux certs' attrs ()) + +-- | Rewrite the bindings introduced by an expression (to ensure they are +-- unique) and fix any dependencies that are broken as a result of migration or +-- rewriting. +rewriteExp :: Exp GPU -> RewriteM (Exp GPU) +rewriteExp = + mapExpM $ + Mapper + { mapOnSubExp = renameSubExp, + mapOnBody = const rewriteBody, + mapOnVName = rename, + mapOnRetType = renameExtType, + mapOnBranchType = renameExtType, + mapOnFParam = rewriteParam, + mapOnLParam = rewriteParam, + mapOnOp = const opError + } + where + -- This indicates that something fundamentally is wrong with the migration + -- table produced by the "Futhark.Analysis.MigrationTable" module. + opError = compilerBugS "Cannot migrate a host-only operation to device." + +-- | Rewrite the binding introduced by a single parameter (to ensure it is +-- unique) and fix any dependencies that are broken as a result of migration or +-- rewriting. +rewriteParam :: Param (TypeBase Shape u) -> RewriteM (Param (TypeBase Shape u)) +rewriteParam (Param attrs n t) = do + n' <- rewriteName n + t' <- renameType t + pure (Param attrs n' t') + +-- | Return the name to use for a rewritten dependency. +rename :: VName -> RewriteM VName +rename n = do + st <- get + let renames = rewriteRenames st + let idx = baseTag n + case IM.lookup idx renames of + Just n' -> pure n' + _ -> + do + let stms = rewritePrologue st + (stms', n') <- lift $ useScalar stms n + modify $ \st' -> + st' + { rewriteRenames = IM.insert idx n' renames, + rewritePrologue = stms' + } + pure n' + +-- | Update the variable names within a 'Result' to account for migration and +-- rewriting. +renameResult :: Result -> RewriteM Result +renameResult = mapM renameSubExpRes + +-- | Update the variable names within a 'SubExpRes' to account for migration and +-- rewriting. +renameSubExpRes :: SubExpRes -> RewriteM SubExpRes +renameSubExpRes (SubExpRes certs se) = do + certs' <- renameCerts certs + se' <- renameSubExp se + pure (SubExpRes certs' se') + +-- | Update the variable names of certificates to account for migration and +-- rewriting. +renameCerts :: Certs -> RewriteM Certs +renameCerts cs = Certs <$> mapM rename (unCerts cs) + +-- | Update any variable name within a t'SubExp' to account for migration and +-- rewriting. +renameSubExp :: SubExp -> RewriteM SubExp +renameSubExp (Var n) = Var <$> rename n +renameSubExp se = pure se + +-- | Update the variable names within a type to account for migration and +-- rewriting. +renameType :: TypeBase Shape u -> RewriteM (TypeBase Shape u) +-- Note: mapOnType also maps the VName token of accumulators +renameType = mapOnType renameSubExp + +-- | Update the variable names within an existential type to account for +-- migration and rewriting. +renameExtType :: TypeBase ExtShape u -> RewriteM (TypeBase ExtShape u) +-- Note: mapOnExtType also maps the VName token of accumulators +renameExtType = mapOnExtType renameSubExp diff --git a/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs b/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable.hs @@ -0,0 +1,1620 @@ +-- | +-- This module implements program analysis to determine which program statements +-- the "Futhark.Optimise.ReduceDeviceSyncs" pass should move into 'GPUBody' kernels +-- to reduce blocking memory transfers between host and device. The results of +-- the analysis is encoded into a 'MigrationTable' which can be queried. +-- +-- To reduce blocking scalar reads the module constructs a data flow +-- dependency graph of program variables (see +-- "Futhark.Optimise.ReduceDeviceSyncs.MigrationTable.Graph") in which +-- it finds a minimum vertex cut that separates array reads of scalars +-- from transitive usage that cannot or should not be migrated to +-- device. +-- +-- The variables of each partition are assigned a 'MigrationStatus' that states +-- whether the computation of those variables should be moved to device or +-- remain on host. Due to how the graph is built and the vertex cut is found all +-- variables bound by a single statement will belong to the same partition. +-- +-- The vertex cut contains all variables that will reside in device memory but +-- are required by host operations. These variables must be read from device +-- memory and cannot be reduced further in number merely by migrating +-- statements (subject to the accuracy of the graph model). The model is built +-- to reduce the worst-case number of scalar reads; an optimal migration of +-- statements depends on runtime data. +-- +-- Blocking scalar writes are reduced by either turning such writes into +-- asynchronous kernels, as is done with scalar array literals and accumulator +-- updates, or by transforming host-device writing into device-device copying. +-- +-- For details on how the graph is constructed and how the vertex cut is found, +-- see the master thesis "Reducing Synchronous GPU Memory Transfers" by Philip +-- Børgesen (2022). +module Futhark.Optimise.ReduceDeviceSyncs.MigrationTable + ( -- * Analysis + analyseProg, + + -- * Types + MigrationTable, + MigrationStatus (..), + + -- * Query + + -- | These functions all assume that no parent statement should be migrated. + -- That is @shouldMoveStm stm mt@ should return @False@ for every statement + -- @stm@ with a body that a queried 'VName' or 'Stm' is nested within, + -- otherwise the query result may be invalid. + shouldMoveStm, + shouldMove, + usedOnHost, + statusOf, + ) +where + +import Control.Monad +import Control.Monad.Trans.Class +import qualified Control.Monad.Trans.Reader as R +import Control.Monad.Trans.State.Strict () +import Control.Monad.Trans.State.Strict hiding (State) +import Control.Parallel.Strategies (parMap, rpar) +import Data.Bifunctor (first, second) +import Data.Foldable +import qualified Data.IntMap.Strict as IM +import qualified Data.IntSet as IS +import qualified Data.Map.Strict as M +import Data.Maybe (fromJust, fromMaybe, isJust, isNothing) +import qualified Data.Sequence as SQ +import Data.Set (Set, (\\)) +import qualified Data.Set as S +import Futhark.Error +import Futhark.IR.GPU +import Futhark.Optimise.ReduceDeviceSyncs.MigrationTable.Graph + ( EdgeType (..), + Edges (..), + Id, + IdSet, + Result (..), + Routing (..), + Vertex (..), + ) +import qualified Futhark.Optimise.ReduceDeviceSyncs.MigrationTable.Graph as MG + +-------------------------------------------------------------------------------- +-- MIGRATION TABLES -- +-------------------------------------------------------------------------------- + +-- | Where the value bound by a name should be computed. +data MigrationStatus + = -- | The statement that computes the value should be moved to device. + -- No host usage of the value will be left after the migration. + MoveToDevice + | -- | As 'MoveToDevice' but host usage of the value will remain after + -- migration. + UsedOnHost + | -- | The statement that computes the value should remain on host. + StayOnHost + deriving (Eq, Ord, Show) + +-- | Identifies +-- +-- (1) which statements should be moved from host to device to reduce the +-- worst case number of blocking memory transfers. +-- +-- (2) which migrated variables that still will be used on the host after +-- all such statements have been moved. +newtype MigrationTable = MigrationTable (IM.IntMap MigrationStatus) + +-- | Where should the value bound by this name be computed? +statusOf :: VName -> MigrationTable -> MigrationStatus +statusOf n (MigrationTable mt) = + fromMaybe StayOnHost $ IM.lookup (baseTag n) mt + +-- | Should this whole statement be moved from host to device? +shouldMoveStm :: Stm GPU -> MigrationTable -> Bool +shouldMoveStm (Let (Pat ((PatElem n _) : _)) _ (BasicOp (Index _ slice))) mt = + statusOf n mt == MoveToDevice || any movedOperand slice + where + movedOperand (Var op) = statusOf op mt == MoveToDevice + movedOperand _ = False +shouldMoveStm (Let (Pat ((PatElem n _) : _)) _ (BasicOp _)) mt = + statusOf n mt /= StayOnHost +shouldMoveStm (Let (Pat ((PatElem n _) : _)) _ Apply {}) mt = + statusOf n mt /= StayOnHost +shouldMoveStm (Let _ _ (If (Var n) _ _ _)) mt = + statusOf n mt == MoveToDevice +shouldMoveStm (Let _ _ (DoLoop _ (ForLoop _ _ (Var n) _) _)) mt = + statusOf n mt == MoveToDevice +shouldMoveStm (Let _ _ (DoLoop _ (WhileLoop n) _)) mt = + statusOf n mt == MoveToDevice +-- BasicOp and Apply statements might not bind any variables (shouldn't happen). +-- If statements might use a constant branch condition. +-- For loop statements might use a constant number of iterations. +-- HostOp statements cannot execute on device. +-- WithAcc statements are never moved in their entirety. +shouldMoveStm _ _ = False + +-- | Should the value bound by this name be computed on device? +shouldMove :: VName -> MigrationTable -> Bool +shouldMove n mt = statusOf n mt /= StayOnHost + +-- | Will the value bound by this name be used on host? +usedOnHost :: VName -> MigrationTable -> Bool +usedOnHost n mt = statusOf n mt /= MoveToDevice + +-- | Merges two migration tables that are assumed to be disjoint. +merge :: MigrationTable -> MigrationTable -> MigrationTable +merge (MigrationTable a) (MigrationTable b) = MigrationTable (a `IM.union` b) + +-------------------------------------------------------------------------------- +-- HOST-ONLY FUNCTION ANALYSIS -- +-------------------------------------------------------------------------------- + +-- | Identifies top-level function definitions that cannot be run on the +-- device. The application of any such function is host-only. +type HostOnlyFuns = Set Name + +-- | Returns the names of all top-level functions that cannot be called from the +-- device. The evaluation of such a function is host-only. +hostOnlyFunDefs :: [FunDef GPU] -> HostOnlyFuns +hostOnlyFunDefs funs = + let names = map funDefName funs + call_map = M.fromList $ zip names (map checkFunDef funs) + in S.fromList names \\ keysToSet (removeHostOnly call_map) + where + keysToSet = S.fromAscList . M.keys + + removeHostOnly cm = + let (host_only, cm') = M.partition isHostOnly cm + in if M.null host_only + then cm' + else removeHostOnly $ M.map (checkCalls $ keysToSet host_only) cm' + + isHostOnly = isNothing + + -- A function that calls a host-only function is itself host-only. + checkCalls hostOnlyFuns (Just calls) + | hostOnlyFuns `S.disjoint` calls = + Just calls + checkCalls _ _ = + Nothing + +-- | 'checkFunDef' returns 'Nothing' if this function definition uses arrays or +-- HostOps. Otherwise it returns the names of all applied functions, which may +-- include user defined functions that could turn out to be host-only. +checkFunDef :: FunDef GPU -> Maybe (Set Name) +checkFunDef fun = do + checkFParams (funDefParams fun) + checkRetTypes (funDefRetType fun) + checkBody (funDefBody fun) + where + hostOnly = Nothing + ok = Just () + check isArr as = if any isArr as then hostOnly else ok + + checkFParams = check isArray + + checkLParams = check (isArray . fst) + + checkRetTypes = check isArrayType + + checkPats = check isArray + + checkLoopForm (ForLoop _ _ _ (_ : _)) = hostOnly + checkLoopForm _ = ok + + checkBody = checkStms . bodyStms + + checkStms stms = S.unions <$> mapM checkStm stms + + checkStm (Let (Pat pats) _ e) = checkPats pats >> checkExp e + + -- Any expression that produces an array is caught by checkPats + checkExp (BasicOp (Index _ _)) = hostOnly + checkExp (WithAcc _ _) = hostOnly + checkExp (Op _) = hostOnly + checkExp (Apply fn _ _ _) = Just (S.singleton fn) + checkExp (If _ tbranch fbranch _) = do + calls1 <- checkBody tbranch + calls2 <- checkBody fbranch + pure (calls1 <> calls2) + checkExp (DoLoop params lform body) = do + checkLParams params + checkLoopForm lform + checkBody body + checkExp _ = Just S.empty + +-------------------------------------------------------------------------------- +-- MIGRATION ANALYSIS -- +-------------------------------------------------------------------------------- + +-- | HostUsage identifies scalar variables that are used on host. +type HostUsage = [Id] + +nameToId :: VName -> Id +nameToId = baseTag + +-- | Analyses a program to return a migration table that covers all its +-- statements and variables. +analyseProg :: Prog GPU -> MigrationTable +analyseProg (Prog consts funs) = + let hof = hostOnlyFunDefs funs + mt = analyseConsts hof funs consts + mts = parMap rpar (analyseFunDef hof) funs + in foldl' merge mt mts + +-- | Analyses top-level constants. +analyseConsts :: HostOnlyFuns -> [FunDef GPU] -> Stms GPU -> MigrationTable +analyseConsts hof funs consts = + let usage = M.foldlWithKey (f $ freeIn funs) [] (scopeOf consts) + in analyseStms hof usage consts + where + f free usage n t + | isScalar t, + n `nameIn` free = + nameToId n : usage + | otherwise = + usage + +-- | Analyses a top-level function definition. +analyseFunDef :: HostOnlyFuns -> FunDef GPU -> MigrationTable +analyseFunDef hof fd = + let body = funDefBody fd + usage = foldl' f [] $ zip (bodyResult body) (funDefRetType fd) + stms = bodyStms body + in analyseStms hof usage stms + where + f usage (SubExpRes _ (Var n), t) | isScalarType t = nameToId n : usage + f usage _ = usage + +-- | Analyses statements. The 'HostUsage' list identifies which bound scalar +-- variables that subsequently may be used on host. All free variables such as +-- constants and function parameters are assumed to reside on host. +analyseStms :: HostOnlyFuns -> HostUsage -> Stms GPU -> MigrationTable +analyseStms hof usage stms = + let (g, srcs, _) = buildGraph hof usage stms + (routed, unrouted) = srcs + (_, g') = MG.routeMany unrouted g -- hereby routed + f st' = MG.fold g' visit st' Normal + st = foldl' f (initial, MG.none) unrouted + (vr, vn, tn) = fst $ foldl' f st routed + in -- TODO: Delay reads into (deeper) branches + + MigrationTable $ + IM.unions + [ IM.fromSet (const MoveToDevice) vr, + IM.fromSet (const MoveToDevice) vn, + -- Read by host if not reached by a reversed edge + IM.fromSet (const UsedOnHost) tn + ] + where + -- 1) Visited by reversed edge. + -- 2) Visited by normal edge, no route. + -- 3) Visited by normal edge, had route; will potentially be read by host. + initial = (IS.empty, IS.empty, IS.empty) + + visit (vr, vn, tn) Reversed v = + let vr' = IS.insert (vertexId v) vr + in (vr', vn, tn) + visit (vr, vn, tn) Normal v@Vertex {vertexRouting = NoRoute} = + let vn' = IS.insert (vertexId v) vn + in (vr, vn', tn) + visit (vr, vn, tn) Normal v = + let tn' = IS.insert (vertexId v) tn + in (vr, vn, tn') + +-------------------------------------------------------------------------------- +-- TYPE HELPERS -- +-------------------------------------------------------------------------------- + +isScalar :: Typed t => t -> Bool +isScalar = isScalarType . typeOf + +isScalarType :: TypeBase shape u -> Bool +isScalarType (Prim Unit) = False +isScalarType (Prim _) = True +isScalarType _ = False + +isArray :: Typed t => t -> Bool +isArray = isArrayType . typeOf + +isArrayType :: ArrayShape shape => TypeBase shape u -> Bool +isArrayType = (0 <) . arrayRank + +-------------------------------------------------------------------------------- +-- GRAPH BUILDING -- +-------------------------------------------------------------------------------- + +buildGraph :: HostOnlyFuns -> HostUsage -> Stms GPU -> (Graph, Sources, Sinks) +buildGraph hof usage stms = + let (g, srcs, sinks) = execGrapher hof (graphStms stms) + g' = foldl' (flip MG.connectToSink) g usage + in (g', srcs, sinks) + +-- | Graph a body. +graphBody :: Body GPU -> Grapher () +graphBody body = do + let res_ops = namesIntSet $ freeIn (bodyResult body) + body_stats <- + captureBodyStats $ + incBodyDepthFor (graphStms (bodyStms body) >> tellOperands res_ops) + + body_depth <- (1 +) <$> getBodyDepth + let host_only = IS.member body_depth (bodyHostOnlyParents body_stats) + modify $ \st -> + let stats = stateStats st + hops' = IS.delete body_depth (bodyHostOnlyParents stats) + -- If body contains a variable that is required on host the parent + -- statement that contains this body cannot be migrated as a whole. + stats' = if host_only then stats {bodyHostOnly = True} else stats + in st {stateStats = stats' {bodyHostOnlyParents = hops'}} + +-- | Graph multiple statements. +graphStms :: Stms GPU -> Grapher () +graphStms = mapM_ graphStm + +-- | Graph a single statement. +graphStm :: Stm GPU -> Grapher () +graphStm stm = do + let bs = boundBy stm + let e = stmExp stm + -- IMPORTANT! It is generally assumed that all scalars within types and + -- shapes are present on host. Any expression of a type wherein one of its + -- scalar operands appears must therefore ensure that that scalar operand is + -- marked as a size variable (see the 'hostSize' function). + case e of + BasicOp (SubExp se) -> do + graphSimple bs e + one bs `reusesSubExp` se + BasicOp (Opaque _ se) -> do + graphSimple bs e + one bs `reusesSubExp` se + BasicOp (ArrayLit arr t) + | isScalar t, + any (isJust . subExpVar) arr -> + -- Migrating an array literal with free variables saves a write for + -- every scalar it contains. Under some backends the compiler + -- generates asynchronous writes for scalar constants but otherwise + -- each write will be synchronous. If all scalars are constants then + -- the compiler generates more efficient code that copies static + -- device memory. + graphAutoMove (one bs) + BasicOp UnOp {} -> graphSimple bs e + BasicOp BinOp {} -> graphSimple bs e + BasicOp CmpOp {} -> graphSimple bs e + BasicOp ConvOp {} -> graphSimple bs e + BasicOp Assert {} -> + -- == OpenCL ============================================================= + -- + -- The next read after the execution of a kernel containing an assertion + -- will be made asynchronous, followed by an asynchronous read to check + -- if any assertion failed. The runtime will then block for all enqueued + -- operations to finish. + -- + -- Since an assertion only binds a certificate of unit type, an assertion + -- cannot increase the number of (read) synchronizations that occur. In + -- this regard it is free to migrate. The synchronization that does occur + -- is however (presumably) more expensive as the pipeline of GPU work will + -- be flushed. + -- + -- Since this cost is difficult to quantify and amortize over assertion + -- migration candidates (cost depends on ordering of kernels and reads) we + -- assume it is insignificant. This will likely hold for a system where + -- multiple threads or processes schedules GPU work, as system-wide + -- throughput only will decrease if the GPU utilization decreases as a + -- result. + -- + -- == CUDA =============================================================== + -- + -- Under the CUDA backend every read is synchronous and is followed by + -- a full synchronization that blocks for all enqueued operations to + -- finish. If any enqueued kernel contained an assertion, another + -- synchronous read is then made to check if an assertion failed. + -- + -- Migrating an assertion to save a read may thus introduce new reads, and + -- the total number of reads can hence either decrease, remain the same, + -- or even increase, subject to the ordering of reads and kernels that + -- perform assertions. + -- + -- Since it is possible to implement the same failure checking scheme as + -- OpenCL using asynchronous reads (and doing so would be a good idea!) + -- we consider this to be acceptable. + -- + -- TODO: Implement the OpenCL failure checking scheme under CUDA. This + -- should reduce the number of synchronizations per read to one. + graphSimple bs e + BasicOp (Index _ slice) + | isFixed slice -> + graphRead (one bs) + BasicOp {} + | [(_, t)] <- bs, + dims <- arrayDims t, + dims /= [], -- i.e. produces an array + all (== intConst Int64 1) dims -> + -- An expression that produces an array that only contains a single + -- primitive value is as efficient to compute and copy as a scalar, + -- and introduces no size variables. + -- + -- This is an exception to the inefficiency rules that comes next. + graphSimple bs e + -- Expressions with a cost sublinear to the size of their result arrays are + -- risky to migrate as we cannot guarantee that their results are not + -- returned from a GPUBody, which always copies its return values. Since + -- this would make the effective asymptotic cost of such statements linear + -- we block them from being migrated on their own. + -- + -- The parent statement of an enclosing body may still be migrated as a + -- whole given that each of its returned arrays either + -- 1) is backed by memory used by a migratable statement within its body. + -- 2) contains just a single element. + -- An array matching either criterion is denoted "copyable memory" because + -- the asymptotic cost of copying it is less than or equal to the statement + -- that produced it. This makes the parent of statements with sublinear cost + -- safe to migrate. + BasicOp (Index arr s) -> do + graphInefficientReturn (sliceDims s) e + one bs `reuses` arr + BasicOp (Update _ arr _ _) -> do + graphInefficientReturn [] e + one bs `reuses` arr + BasicOp (FlatIndex arr s) -> do + -- Migrating a FlatIndex leads to a memory allocation error. + -- + -- TODO: Fix FlatIndex memory allocation error. + -- + -- Can be replaced with 'graphHostOnly e' to disable migration. + -- A fix can be verified by enabling tests/migration/reuse2_flatindex.fut + graphInefficientReturn (flatSliceDims s) e + one bs `reuses` arr + BasicOp (FlatUpdate arr _ _) -> do + graphInefficientReturn [] e + one bs `reuses` arr + BasicOp (Scratch _ s) -> + -- Migrating a Scratch leads to a memory allocation error. + -- + -- TODO: Fix Scratch memory allocation error. + -- + -- Can be replaced with 'graphHostOnly e' to disable migration. + -- A fix can be verified by enabling tests/migration/reuse4_scratch.fut + graphInefficientReturn s e + BasicOp (Reshape s arr) -> do + graphInefficientReturn (newDims s) e + one bs `reuses` arr + BasicOp (Rearrange _ arr) -> do + graphInefficientReturn [] e + one bs `reuses` arr + BasicOp (Rotate _ arr) -> do + -- Migrating a Rotate leads to a memory allocation error. + -- + -- TODO: Fix Rotate memory allocation error. + -- + -- Can be replaced with 'graphHostOnly e' to disable migration. + -- A fix can be verified by enabling tests/migration/reuse7_rotate.fut + graphInefficientReturn [] e + one bs `reuses` arr + -- Expressions with a cost linear to the size of their result arrays are + -- inefficient to migrate into GPUBody kernels as such kernels are single- + -- threaded. For sufficiently large arrays the cost may exceed what is saved + -- by avoiding reads. We therefore also block these from being migrated, + -- as well as their parents. + BasicOp ArrayLit {} -> + -- An array literal purely of primitive constants can be hoisted out to be + -- a top-level constant, unless it is to be returned or consumed. + -- Otherwise its runtime implementation will copy a precomputed static + -- array and thus behave like a 'Copy'. + -- Whether the rows are primitive constants or arrays, without any scalar + -- variable operands such ArrayLit cannot directly prevent a scalar read. + graphHostOnly e + BasicOp Concat {} -> + -- Is unlikely to prevent a scalar read as the only SubExp operand in + -- practice is a computation of host-only size variables. + graphHostOnly e + BasicOp Copy {} -> + -- Only takes an array operand, so cannot directly prevent a scalar read. + graphHostOnly e + BasicOp Manifest {} -> + -- Takes no scalar operands so cannot directly prevent a scalar read. + -- It is introduced as part of the BlkRegTiling kernel optimization and + -- is thus unlikely to prevent the migration of a parent which was not + -- already blocked by some host-only operation. + graphHostOnly e + BasicOp Iota {} -> graphHostOnly e + BasicOp Replicate {} -> graphHostOnly e + -- END + BasicOp UpdateAcc {} -> + graphUpdateAcc (one bs) e + Apply fn _ _ _ -> + graphApply fn bs e + If cond tbody fbody _ -> + graphIf bs cond tbody fbody + DoLoop params lform body -> + graphLoop bs params lform body + WithAcc inputs f -> + graphWithAcc bs inputs f + Op GPUBody {} -> + -- A GPUBody can be migrated into a parent GPUBody by replacing it with + -- its body statements and binding its return values inside 'ArrayLit's. + tellGPUBody + Op _ -> + graphHostOnly e + where + one [x] = x + one _ = compilerBugS "Type error: unexpected number of pattern elements." + + isFixed = isJust . sliceIndices + + -- new_dims may introduce new size variables which must be present on host + -- when this expression is evaluated. + graphInefficientReturn new_dims e = do + mapM_ hostSize new_dims + graphedScalarOperands e >>= addEdges ToSink + + hostSize (Var n) = hostSizeVar n + hostSize _ = pure () + + hostSizeVar = requiredOnHost . nameToId + +-- | Bindings for all pattern elements bound by a statement. +boundBy :: Stm GPU -> [Binding] +boundBy = map (\(PatElem n t) -> (nameToId n, t)) . patElems . stmPat + +-- | Graph a statement which in itself neither reads scalars from device memory +-- nor forces such scalars to be available on host. Such statement can be moved +-- to device to eliminate the host usage of its operands which transitively may +-- depend on a scalar device read. +graphSimple :: [Binding] -> Exp GPU -> Grapher () +graphSimple bs e = do + -- Only add vertices to the graph if they have a transitive dependency to + -- an array read. Transitive dependencies through variables connected to + -- sinks do not count. + ops <- graphedScalarOperands e + let edges = MG.declareEdges (map fst bs) + unless (IS.null ops) (mapM_ addVertex bs >> addEdges edges ops) + +-- | Graph a statement that reads a scalar from device memory. +graphRead :: Binding -> Grapher () +graphRead b = do + -- Operands are not important as the source will block routes through b. + addSource b + tellRead + +-- | Graph a statement that always should be moved to device. +graphAutoMove :: Binding -> Grapher () +graphAutoMove = + -- Operands are not important as the source will block routes through b. + addSource + +-- | Graph a statement that is unfit for execution in a GPUBody and thus must +-- be executed on host, requiring all its operands to be made available there. +-- Parent statements of enclosing bodies are also blocked from being migrated. +graphHostOnly :: Exp GPU -> Grapher () +graphHostOnly e = do + -- Connect the vertices of all operands to sinks to mark that they are + -- required on host. Transitive reads that they depend upon can be delayed + -- no further, and any parent statements cannot be migrated. + ops <- graphedScalarOperands e + addEdges ToSink ops + tellHostOnly + +-- | Graph an 'UpdateAcc' statement. +graphUpdateAcc :: Binding -> Exp GPU -> Grapher () +graphUpdateAcc b e | (_, Acc a _ _ _) <- b = + -- The actual graphing is delayed to the corrensponding 'WithAcc' parent. + modify $ \st -> + let accs = stateUpdateAccs st + accs' = IM.alter add (nameToId a) accs + in st {stateUpdateAccs = accs'} + where + add Nothing = Just [(b, e)] + add (Just xs) = Just $ (b, e) : xs +graphUpdateAcc _ _ = + compilerBugS + "Type error: UpdateAcc did not produce accumulator typed value." + +-- | Graph a function application. +graphApply :: Name -> [Binding] -> Exp GPU -> Grapher () +graphApply fn bs e = do + hof <- isHostOnlyFun fn + if hof + then graphHostOnly e + else graphSimple bs e + +-- | Graph an if statement. +graphIf :: [Binding] -> SubExp -> Body GPU -> Body GPU -> Grapher () +graphIf bs cond tbody fbody = do + body_host_only <- + incForkDepthFor + ( do + tstats <- captureBodyStats (graphBody tbody) + fstats <- captureBodyStats (graphBody fbody) + pure $ bodyHostOnly tstats || bodyHostOnly fstats + ) + + -- Record aliases for copyable memory backing returned arrays. + may_copy_results <- reusesBranches bs (results tbody) (results fbody) + let may_migrate = not body_host_only && may_copy_results + + cond_id <- case (may_migrate, cond) of + (False, Var n) -> + -- The migration status of the condition is what determines whether the + -- statement may be migrated as a whole or not. See 'shouldMoveStm'. + connectToSink (nameToId n) >> pure IS.empty + (True, Var n) -> onlyGraphedScalar n + (_, _) -> pure IS.empty + + tellOperands cond_id + + -- Connect branch results to bound variables to allow delaying reads out of + -- branches. It might also be beneficial to move the whole statement to + -- device, to avoid reading the branch condition value. This must be balanced + -- against the need to read the values bound by the if statement. + -- + -- By connecting the branch condition to each variable bound by the statement + -- the condition will only stay on device if + -- + -- (1) the if statement is not required on host, based on the statements + -- within its body. + -- + -- (2) no additional reads will be required to use the if statement bound + -- variables should the whole statement be migrated. + -- + -- If the condition is migrated to device and stays there, then the if + -- statement must necessarily execute on device. + -- + -- While the graph model built by this module generally migrates no more + -- statements than necessary to obtain a minimum vertex cut, the branches + -- of if statements are subject to an inaccuracy. Specifically model is not + -- strong enough to capture their mutual exclusivity and thus encodes that + -- both branches are taken. While this does not affect the resulting number + -- of host-device reads it means that some reads may needlessly be delayed + -- out of branches. The overhead as measured on futhark-benchmarks appears + -- to be neglible though. + ret <- zipWithM (comb cond_id) (bodyResult tbody) (bodyResult fbody) + mapM_ (uncurry createNode) (zip bs ret) + where + results = map resSubExp . bodyResult + + comb ci a b = (ci <>) <$> onlyGraphedScalars (toSet a <> toSet b) + + toSet (SubExpRes _ (Var n)) = S.singleton n + toSet _ = S.empty + +----------------------------------------------------- +-- These type aliases are only used by 'graphLoop' -- +----------------------------------------------------- +type ReachableBindings = IdSet + +type ReachableBindingsCache = MG.Visited (MG.Result ReachableBindings) + +type NonExhausted = [Id] + +type LoopValue = (Binding, Id, SubExp, SubExp) + +----------------------------------------------------- +----------------------------------------------------- + +-- | Graph a loop statement. +graphLoop :: + [Binding] -> + [(FParam GPU, SubExp)] -> + LoopForm GPU -> + Body GPU -> + Grapher () +graphLoop [] _ _ _ = + -- We expect each loop to bind a value or be eliminated. + compilerBugS "Loop statement bound no variable; should have been eliminated." +graphLoop (b : bs) params lform body = do + -- Graph loop params and body while capturing statistics. + g0 <- getGraph + stats <- captureBodyStats (subgraphId `graphIdFor` graphTheLoop) + + -- Record aliases for copyable memory backing returned arrays. + -- Does the loop return any arrays which prevent it from being migrated? + let args = map snd params + let results = map resSubExp (bodyResult body) + may_copy_results <- reusesBranches (b : bs) args results + + -- Connect loop condition to a sink if the loop cannot be migrated. + -- The migration status of the condition is what determines whether the + -- loop may be migrated as a whole or not. See 'shouldMoveStm'. + let may_migrate = not (bodyHostOnly stats) && may_copy_results + unless may_migrate $ case lform of + ForLoop _ _ (Var n) _ -> connectToSink (nameToId n) + WhileLoop n -> connectToSink (nameToId n) + _ -> pure () + + -- Connect graphed return values to their loop parameters. + g1 <- getGraph + mapM_ (mergeLoopParam g1) loopValues + + -- Route the sources within the loop body in isolation. + -- The loop graph must not be altered after this point. + srcs <- routeSubgraph subgraphId + + -- Graph the variables bound by the statement. + forM_ loopValues $ \(bnd, p, _, _) -> createNode bnd (IS.singleton p) + + -- If a device read is delayed from one iteration to the next the + -- corresponding variables bound by the statement must be treated as + -- sources. + g2 <- getGraph + let (dbs, rbc) = foldl' (deviceBindings g2) (IS.empty, MG.none) srcs + modifySources $ second (IS.toList dbs <>) + + -- Connect operands to sinks if they can reach a sink within the loop. + -- Otherwise connect them to the loop bound variables that they can + -- reach and exhaust their normal entry edges into the loop. + -- This means a read can be delayed through a loop but not into it if + -- that would increase the number of reads done by any given iteration. + let ops = IS.filter (`MG.member` g0) (bodyOperands stats) + foldM_ connectOperand rbc (IS.elems ops) + + -- It might be beneficial to move the whole loop to device, to avoid + -- reading the (initial) loop condition value. This must be balanced + -- against the need to read the values bound by the loop statement. + -- + -- For more details see the similar description for if statements. + when may_migrate $ case lform of + ForLoop _ _ n _ -> + onlyGraphedScalarSubExp n >>= addEdges (ToNodes bindings Nothing) + WhileLoop n + | (_, _, arg, _) <- loopValueFor n -> + onlyGraphedScalarSubExp arg >>= addEdges (ToNodes bindings Nothing) + where + subgraphId :: Id + subgraphId = fst b + + loopValues :: [LoopValue] + loopValues = + let tmp = zip3 (b : bs) params (bodyResult body) + tmp' = flip map tmp $ + \(bnd, (p, arg), res) -> + let i = nameToId (paramName p) + in (bnd, i, arg, resSubExp res) + in filter (\((_, t), _, _, _) -> isScalar t) tmp' + + bindings :: IdSet + bindings = IS.fromList $ map (\((i, _), _, _, _) -> i) loopValues + + loopValueFor :: VName -> LoopValue + loopValueFor n = + fromJust $ find (\(_, p, _, _) -> p == nameToId n) loopValues + + graphTheLoop :: Grapher () + graphTheLoop = do + mapM_ graphParam loopValues + + -- For simplicity we do not currently track memory reuse through merge + -- parameters. A parameter does not simply reuse the memory of its + -- argument; it must also consider the iteration return value, which in + -- turn may depend on other merge parameters. + -- + -- Situations that would benefit from this tracking is unlikely to occur + -- at the time of writing, and if it occurs current compiler limitations + -- will prevent successful compilation. + -- Specifically it requires the merge parameter argument to reuse memory + -- from an array literal, and both it and the loop must occur within an + -- if statement branch. Array literals are generally hoisted out of if + -- statements however, and when they are not, a memory allocation error + -- occurs. + -- + -- TODO: Track memory reuse through merge parameters. + + case lform of + ForLoop _ _ n elems -> do + onlyGraphedScalarSubExp n >>= tellOperands + mapM_ graphForInElem elems + WhileLoop _ -> pure () + graphBody body + where + graphForInElem (p, arr) = do + when (isScalar p) $ addSource (nameToId $ paramName p, typeOf p) + when (isArray p) $ (nameToId (paramName p), typeOf p) `reuses` arr + + graphParam ((_, t), p, arg, _) = + do + -- It is unknown whether a read can be delayed via the parameter + -- from one iteration to the next, so we have to create a vertex + -- even if the initial value never depends on a read. + addVertex (p, t) + ops <- onlyGraphedScalarSubExp arg + addEdges (MG.oneEdge p) ops + + mergeLoopParam :: Graph -> LoopValue -> Grapher () + mergeLoopParam g (_, p, _, res) + | Var n <- res, + ret <- nameToId n, + ret /= p = + if MG.isSinkConnected p g + then connectToSink ret + else addEdges (MG.oneEdge p) (IS.singleton ret) + | otherwise = + pure () + deviceBindings :: + Graph -> + (ReachableBindings, ReachableBindingsCache) -> + Id -> + (ReachableBindings, ReachableBindingsCache) + deviceBindings g (rb, rbc) i = + let (r, rbc') = MG.reduce g bindingReach rbc Normal i + in case r of + Produced rb' -> (rb <> rb', rbc') + _ -> + compilerBugS + "Migration graph sink could be reached from source after it\ + \ had been attempted routed." + bindingReach :: + ReachableBindings -> + EdgeType -> + Vertex Meta -> + ReachableBindings + bindingReach rb _ v + | i <- vertexId v, + IS.member i bindings = + IS.insert i rb + | otherwise = + rb + connectOperand :: + ReachableBindingsCache -> + Id -> + Grapher ReachableBindingsCache + connectOperand cache op = do + g <- getGraph + case MG.lookup op g of + Nothing -> pure cache + Just v -> + case vertexEdges v of + ToSink -> pure cache + ToNodes es Nothing -> connectOp g cache op es + ToNodes _ (Just nx) -> connectOp g cache op nx + where + connectOp :: + Graph -> + ReachableBindingsCache -> + Id -> -- operand id + IdSet -> -- its edges + Grapher ReachableBindingsCache + connectOp g rbc i es = do + let (res, nx, rbc') = findBindings g (IS.empty, [], rbc) (IS.elems es) + case res of + FoundSink -> connectToSink i + Produced rb -> modifyGraph $ MG.adjust (updateEdges nx rb) i + pure rbc' + updateEdges :: + NonExhausted -> + ReachableBindings -> + Vertex Meta -> + Vertex Meta + updateEdges nx rb v + | ToNodes es _ <- vertexEdges v = + let nx' = IS.fromList nx + es' = ToNodes (rb <> es) $ Just (rb <> nx') + in v {vertexEdges = es'} + | otherwise = v + findBindings :: + Graph -> + (ReachableBindings, NonExhausted, ReachableBindingsCache) -> + [Id] -> -- current non-exhausted edges + (MG.Result ReachableBindings, NonExhausted, ReachableBindingsCache) + findBindings _ (rb, nx, rbc) [] = + (Produced rb, nx, rbc) + findBindings g (rb, nx, rbc) (i : is) + | Just v <- MG.lookup i g, + Just gid <- metaGraphId (vertexMeta v), + gid == subgraphId -- only search the subgraph + = + let (res, rbc') = MG.reduce g bindingReach rbc Normal i + in case res of + FoundSink -> (FoundSink, [], rbc') + Produced rb' -> findBindings g (rb <> rb', nx, rbc') is + | otherwise = + -- don't exhaust + findBindings g (rb, i : nx, rbc) is + +-- | Graph a 'WithAcc' statement. +graphWithAcc :: + [Binding] -> + [WithAccInput GPU] -> + Lambda GPU -> + Grapher () +graphWithAcc bs inputs f = do + -- Graph the body, capturing 'UpdateAcc' statements for delayed graphing. + graphBody (lambdaBody f) + + -- Graph each accumulator monoid and its associated 'UpdateAcc' statements. + mapM_ graph $ zip (lambdaReturnType f) inputs + + -- Record aliases for the backing memory of each returned array. + -- 'WithAcc' statements are never migrated as a whole and always returns + -- arrays backed by memory allocated elsewhere. + let arrs = concatMap (\(_, as, _) -> map Var as) inputs + let res = drop (length inputs) (bodyResult $ lambdaBody f) + _ <- reusesReturn bs (arrs ++ map resSubExp res) + + -- Connect return variables to bound values. No outgoing edge exists + -- from an accumulator vertex so skip those. Note that accumulators do + -- not map to returned arrays one-to-one but one-to-many. + ret <- mapM (onlyGraphedScalarSubExp . resSubExp) res + mapM_ (uncurry createNode) $ zip (drop (length arrs) bs) ret + where + graph (Acc a _ types _, (_, _, comb)) = do + let i = nameToId a + + delayed <- fromMaybe [] <$> gets (IM.lookup i . stateUpdateAccs) + modify $ \st -> st {stateUpdateAccs = IM.delete i (stateUpdateAccs st)} + + graphAcc i types (fst <$> comb) delayed + + -- Neutral elements must always be made available on host for 'WithAcc' + -- to type check. + mapM_ connectSubExpToSink $ maybe [] snd comb + graph _ = + compilerBugS "Type error: WithAcc expression did not return accumulator." + +-- Graph the operator and all 'UpdateAcc' statements associated with an +-- accumulator. +-- +-- The arguments are the 'Id' for the accumulator token, the element types of +-- the accumulator/operator, its combining function if any, and all associated +-- 'UpdateAcc' statements outside kernels. +graphAcc :: Id -> [Type] -> Maybe (Lambda GPU) -> [Delayed] -> Grapher () +graphAcc i _ _ [] = addSource (i, Prim Unit) -- Only used on device. +graphAcc i types op delayed = do + -- Accumulators are intended for use within SegOps but in principle the AST + -- allows their 'UpdateAcc's to be used outside a kernel. This case handles + -- that unlikely situation. + + env <- ask + st <- get + + -- Collect statistics about the operator statements. + let lambda = fromMaybe (Lambda [] (Body () SQ.empty []) []) op + let m = graphBody (lambdaBody lambda) + let stats = R.runReader (evalStateT (captureBodyStats m) st) env + -- We treat GPUBody kernels as host-only to not bother rewriting them inside + -- operators and to simplify the analysis. They are unlikely to occur anyway. + -- + -- NOTE: Performance may degrade if a GPUBody is replaced with its contents + -- but the containing operator is used on host. + let host_only = bodyHostOnly stats || bodyHasGPUBody stats + + -- op operands are read from arrays and written back so if any of the operands + -- are scalar then a read can be avoided by moving the UpdateAcc usages to + -- device. If the op itself performs scalar reads its UpdateAcc usages should + -- also be moved. + let does_read = bodyReads stats || any isScalar types + + -- Determine which external variables the operator depends upon. + -- 'bodyOperands' cannot be used as it might exclude operands that were + -- connected to sinks within the body, so instead we create an artifical + -- expression to capture graphed operands from. + ops <- graphedScalarOperands (WithAcc [] lambda) + + case (host_only, does_read) of + (True, _) -> do + -- If the operator cannot run well in a GPUBody then all non-kernel + -- UpdateAcc statements are host-only. The current analysis is ignorant + -- of what happens inside kernels so we must assume that the operator + -- is used within a kernel, meaning that we cannot migrate its statements. + -- + -- TODO: Improve analysis if UpdateAcc ever is used outside kernels. + mapM_ (graphHostOnly . snd) delayed + addEdges ToSink ops + (_, True) -> do + -- Migrate all accumulator usage to device to avoid reads and writes. + mapM_ (graphAutoMove . fst) delayed + addSource (i, Prim Unit) + _ -> do + -- Only migrate operator and UpdateAcc statements if it can allow their + -- operands to be migrated. + createNode (i, Prim Unit) ops + forM_ delayed $ + \(b, e) -> graphedScalarOperands e >>= createNode b . IS.insert i + +-- Returns for an expression all scalar operands that must be made available +-- on host to execute the expression there. +graphedScalarOperands :: Exp GPU -> Grapher Operands +graphedScalarOperands e = + let is = fst $ execState (collect e) initial + in IS.intersection is <$> getGraphedScalars + where + initial = (IS.empty, S.empty) -- scalar operands, accumulator tokens + captureName n = modify $ first $ IS.insert (nameToId n) + captureAcc a = modify $ second $ S.insert a + collectFree x = mapM_ captureName (namesToList $ freeIn x) + + collect b@BasicOp {} = + collectBasic b + collect (Apply _ params _ _) = + mapM_ (collectSubExp . fst) params + collect (If cond tbranch fbranch _) = + collectSubExp cond >> collectBody tbranch >> collectBody fbranch + collect (DoLoop params lform body) = do + mapM_ (collectSubExp . snd) params + collectLForm lform + collectBody body + collect (WithAcc accs f) = + collectWithAcc accs f + collect (Op op) = + collectHostOp op + + collectBasic (BasicOp (Update _ _ slice _)) = + -- Writing a scalar to an array can be replaced with copying a single- + -- element slice. If the scalar originates from device memory its read + -- can thus be prevented without requiring the 'Update' to be migrated. + collectFree slice + collectBasic (BasicOp (Replicate shape _)) = + -- The replicate of a scalar can be rewritten as a replicate of a single + -- element array followed by a slice index. + collectFree shape + collectBasic e' = + -- Note: Plain VName values only refer to arrays. + walkExpM (identityWalker {walkOnSubExp = collectSubExp}) e' + + collectSubExp (Var n) = captureName n + collectSubExp _ = pure () + + collectBody body = do + collectStms (bodyStms body) + collectFree (bodyResult body) + collectStms = mapM_ collectStm + + collectStm (Let pat _ ua) + | BasicOp UpdateAcc {} <- ua, + Pat [pe] <- pat, + Acc a _ _ _ <- typeOf pe = + -- Capture the tokens of accumulators used on host. + captureAcc a >> collectBasic ua + collectStm stm = collect (stmExp stm) + + collectLForm (ForLoop _ _ b _) = collectSubExp b + -- WhileLoop condition is declared as a loop parameter. + collectLForm (WhileLoop _) = pure () + + -- The collective operands of an operator lambda body are only used on host + -- if the associated accumulator is used in an UpdateAcc statement outside a + -- kernel. + collectWithAcc inputs f = do + collectBody (lambdaBody f) + used_accs <- gets snd + let accs = take (length inputs) (lambdaReturnType f) + let used = map (\(Acc a _ _ _) -> S.member a used_accs) accs + mapM_ collectAcc (zip used inputs) + + collectAcc (_, (_, _, Nothing)) = pure () + collectAcc (used, (_, _, Just (op, nes))) = do + mapM_ collectSubExp nes + when used $ collectBody (lambdaBody op) + + -- Does not collect named operands in + -- + -- * types and shapes; size variables are assumed available to the host. + -- + -- * use by a kernel body. + -- + -- All other operands are conservatively collected even if they generally + -- appear to be size variables or results computed by a SizeOp. + collectHostOp (SegOp (SegMap lvl sp _ _)) = do + collectSegLevel lvl + collectSegSpace sp + collectHostOp (SegOp (SegRed lvl sp ops _ _)) = do + collectSegLevel lvl + collectSegSpace sp + mapM_ collectSegBinOp ops + collectHostOp (SegOp (SegScan lvl sp ops _ _)) = do + collectSegLevel lvl + collectSegSpace sp + mapM_ collectSegBinOp ops + collectHostOp (SegOp (SegHist lvl sp ops _ _)) = do + collectSegLevel lvl + collectSegSpace sp + mapM_ collectHistOp ops + collectHostOp (SizeOp op) = collectFree op + collectHostOp (OtherOp op) = collectFree op + collectHostOp GPUBody {} = pure () + + collectSegLevel (SegThread (Count num) (Count size) _) = + collectSubExp num >> collectSubExp size + collectSegLevel (SegGroup (Count num) (Count size) _) = + collectSubExp num >> collectSubExp size + + collectSegSpace space = + mapM_ collectSubExp (segSpaceDims space) + + collectSegBinOp (SegBinOp _ _ nes _) = + mapM_ collectSubExp nes + + collectHistOp (HistOp _ rf _ nes _ _) = do + collectSubExp rf + mapM_ collectSubExp nes + +-------------------------------------------------------------------------------- +-- GRAPH BUILDING - PRIMITIVES -- +-------------------------------------------------------------------------------- + +-- | Creates a vertex for the given binding, provided that the set of operands +-- is not empty. +createNode :: Binding -> Operands -> Grapher () +createNode b ops = + unless (IS.null ops) (addVertex b >> addEdges (MG.oneEdge $ fst b) ops) + +-- | Adds a vertex to the graph for the given binding. +addVertex :: Binding -> Grapher () +addVertex (i, t) = do + meta <- getMeta + let v = MG.vertex i meta + when (isScalar t) $ modifyGraphedScalars (IS.insert i) + when (isArray t) $ recordCopyableMemory i (metaBodyDepth meta) + modifyGraph (MG.insert v) + +-- | Adds a source connected vertex to the graph for the given binding. +addSource :: Binding -> Grapher () +addSource b = do + addVertex b + modifySources $ second (fst b :) + +-- | Adds the given edges to each vertex identified by the 'IdSet'. It is +-- assumed that all vertices reside within the body that currently is being +-- graphed. +addEdges :: Edges -> IdSet -> Grapher () +addEdges ToSink is = do + modifyGraph $ \g -> IS.foldl' (flip MG.connectToSink) g is + modifyGraphedScalars (`IS.difference` is) +addEdges es is = do + modifyGraph $ \g -> IS.foldl' (flip $ MG.addEdges es) g is + tellOperands is + +-- | Ensure that a variable (which is in scope) will be made available on host +-- before its first use. +requiredOnHost :: Id -> Grapher () +requiredOnHost i = do + mv <- MG.lookup i <$> getGraph + case mv of + Nothing -> pure () + Just v -> do + connectToSink i + tellHostOnlyParent (metaBodyDepth $ vertexMeta v) + +-- | Connects the vertex of the given id to a sink. +connectToSink :: Id -> Grapher () +connectToSink i = do + modifyGraph (MG.connectToSink i) + modifyGraphedScalars (IS.delete i) + +-- | Like 'connectToSink' but vertex is given by a t'SubExp'. This is a no-op if +-- the t'SubExp' is a constant. +connectSubExpToSink :: SubExp -> Grapher () +connectSubExpToSink (Var n) = connectToSink (nameToId n) +connectSubExpToSink _ = pure () + +-- | Routes all possible routes within the subgraph identified by this id. +-- Returns the ids of the source connected vertices that were attempted routed. +-- +-- Assumption: The subgraph with the given id has just been created and no path +-- exists from it to an external sink. +routeSubgraph :: Id -> Grapher [Id] +routeSubgraph si = do + st <- get + let g = stateGraph st + let (routed, unrouted) = stateSources st + let (gsrcs, unrouted') = span (inSubGraph si g) unrouted + let (sinks, g') = MG.routeMany gsrcs g + put $ + st + { stateGraph = g', + stateSources = (gsrcs ++ routed, unrouted'), + stateSinks = sinks ++ stateSinks st + } + pure gsrcs + +-- | @inSubGraph si g i@ returns whether @g@ contains a vertex with id @i@ that +-- is declared within the subgraph with id @si@. +inSubGraph :: Id -> Graph -> Id -> Bool +inSubGraph si g i + | Just v <- MG.lookup i g, + Just mgi <- metaGraphId (vertexMeta v) = + si == mgi +inSubGraph _ _ _ = False + +-- | @b `reuses` n@ records that @b@ binds an array backed by the same memory +-- as @n@. If @b@ is not array typed or the backing memory is not copyable then +-- this does nothing. +reuses :: Binding -> VName -> Grapher () +reuses (i, t) n + | isArray t = + do + body_depth <- outermostCopyableArray n + case body_depth of + Just bd -> recordCopyableMemory i bd + Nothing -> pure () + | otherwise = + pure () + +reusesSubExp :: Binding -> SubExp -> Grapher () +reusesSubExp b (Var n) = b `reuses` n +reusesSubExp _ _ = pure () + +-- @reusesReturn bs res@ records each array binding in @bs@ as reusing copyable +-- memory if the corresponding return value in @res@ is backed by copyable +-- memory. +-- +-- If every array binding is registered as being backed by copyable memory then +-- the function returns @True@, otherwise it returns @False@. +reusesReturn :: [Binding] -> [SubExp] -> Grapher Bool +reusesReturn bs res = do + body_depth <- metaBodyDepth <$> getMeta + foldM (reuse body_depth) True (zip bs res) + where + reuse :: Int -> Bool -> (Binding, SubExp) -> Grapher Bool + reuse body_depth onlyCopyable (b, se) + | all (== intConst Int64 1) (arrayDims $ snd b) = + -- Single element arrays are immediately recognizable as copyable so + -- don't bother recording those. Note that this case also matches + -- primitive return values. + pure onlyCopyable + | (i, t) <- b, + isArray t, + Var n <- se = + do + res_body_depth <- outermostCopyableArray n + case res_body_depth of + Just inner -> do + recordCopyableMemory i (min body_depth inner) + let returns_free_var = inner <= body_depth + pure (onlyCopyable && not returns_free_var) + _ -> + pure False + | otherwise = + pure onlyCopyable + +-- @reusesBranches bs b1 b2@ records each array binding in @bs@ as reusing +-- copyable memory if each corresponding return value in the lists @b1@ and @b2@ +-- are backed by copyable memory. +-- +-- If every array binding is registered as being backed by copyable memory then +-- the function returns @True@, otherwise it returns @False@. +reusesBranches :: [Binding] -> [SubExp] -> [SubExp] -> Grapher Bool +reusesBranches bs b1 b2 = do + body_depth <- metaBodyDepth <$> getMeta + foldM (reuse body_depth) True $ zip3 bs b1 b2 + where + reuse :: Int -> Bool -> (Binding, SubExp, SubExp) -> Grapher Bool + reuse body_depth onlyCopyable (b, se1, se2) + | all (== intConst Int64 1) (arrayDims $ snd b) = + -- Single element arrays are immediately recognizable as copyable so + -- don't bother recording those. Note that this case also matches + -- primitive return values. + pure onlyCopyable + | (i, t) <- b, + isArray t, + Var n1 <- se1, + Var n2 <- se2 = + do + body_depth_1 <- outermostCopyableArray n1 + body_depth_2 <- outermostCopyableArray n2 + case (body_depth_1, body_depth_2) of + (Just bd1, Just bd2) -> do + let inner = min bd1 bd2 + recordCopyableMemory i (min body_depth inner) + let returns_free_var = inner <= body_depth + pure (onlyCopyable && not returns_free_var) + _ -> + pure False + | otherwise = + pure onlyCopyable + +-------------------------------------------------------------------------------- +-- GRAPH BUILDING - TYPES -- +-------------------------------------------------------------------------------- + +type Grapher = StateT State (R.Reader Env) + +data Env = Env + { -- | See 'HostOnlyFuns'. + envHostOnlyFuns :: HostOnlyFuns, + -- | Metadata for the current body being graphed. + envMeta :: Meta + } + +-- | A measurement of how many bodies something is nested within. +type BodyDepth = Int + +-- | Metadata on the environment that a variable is declared within. +data Meta = Meta + { -- | How many if statement branch bodies the variable binding is nested + -- within. If a route passes through the edge u->v and the fork depth + -- + -- 1) increases from u to v, then u is within a conditional branch. + -- + -- 2) decreases from u to v, then v binds the result of two or more + -- branches. + -- + -- After the graph has been built and routed, this can be used to delay + -- reads into deeper branches to reduce their likelihood of manifesting. + metaForkDepth :: Int, + -- | How many bodies the variable is nested within. + metaBodyDepth :: BodyDepth, + -- | An id for the subgraph within which the variable exists, defined at + -- the body level. A read may only be delayed to a point within its own + -- subgraph. + metaGraphId :: Maybe Id + } + +-- | Ids for all variables used as an operand. +type Operands = IdSet + +-- | Statistics on the statements within a body and their dependencies. +data BodyStats = BodyStats + { -- | Whether the body contained any host-only statements. + bodyHostOnly :: Bool, + -- | Whether the body contained any GPUBody kernels. + bodyHasGPUBody :: Bool, + -- | Whether the body performed any reads. + bodyReads :: Bool, + -- | All scalar variables represented in the graph that have been used + -- as return values of the body or as operands within it, including those + -- that are defined within the body itself. Variables with vertices + -- connected to sinks may be excluded. + bodyOperands :: Operands, + -- | Depth of parent bodies with variables that are required on host. Since + -- the variables are required on host, the parent statements of these bodies + -- cannot be moved to device as a whole. They are host-only. + bodyHostOnlyParents :: IS.IntSet + } + +instance Semigroup BodyStats where + (BodyStats ho1 gb1 r1 o1 hop1) <> (BodyStats ho2 gb2 r2 o2 hop2) = + BodyStats + { bodyHostOnly = ho1 || ho2, + bodyHasGPUBody = gb1 || gb2, + bodyReads = r1 || r2, + bodyOperands = IS.union o1 o2, + bodyHostOnlyParents = IS.union hop1 hop2 + } + +instance Monoid BodyStats where + mempty = + BodyStats + { bodyHostOnly = False, + bodyHasGPUBody = False, + bodyReads = False, + bodyOperands = IS.empty, + bodyHostOnlyParents = IS.empty + } + +type Graph = MG.Graph Meta + +-- | All vertices connected from a source, partitioned into those that have +-- been attempted routed and those which have not. +type Sources = ([Id], [Id]) + +-- | All terminal vertices of routes. +type Sinks = [Id] + +-- | A captured statement for which graphing has been delayed. +type Delayed = (Binding, Exp GPU) + +-- | The vertex handle for a variable and its type. +type Binding = (Id, Type) + +-- | Array variables backed by memory segments that may be copied, mapped to the +-- outermost known body depths that declares arrays backed by a superset of +-- those segments. +type CopyableMemoryMap = IM.IntMap BodyDepth + +data State = State + { -- | The graph being built. + stateGraph :: Graph, + -- | All known scalars that have been graphed. + stateGraphedScalars :: IdSet, + -- | All variables that directly bind scalars read from device memory. + stateSources :: Sources, + -- | Graphed scalars that are used as operands by statements that cannot be + -- migrated. A read cannot be delayed beyond these, so if the statements + -- that bind these variables are moved to device, the variables must be read + -- from device memory. + stateSinks :: Sinks, + -- | Observed 'UpdateAcc' host statements to be graphed later. + stateUpdateAccs :: IM.IntMap [Delayed], + -- | A map of encountered arrays that are backed by copyable memory. + -- Trivial instances such as single element arrays are excluded. + stateCopyableMemory :: CopyableMemoryMap, + -- | Information about the current body being graphed. + stateStats :: BodyStats + } + +-------------------------------------------------------------------------------- +-- GRAPHER OPERATIONS -- +-------------------------------------------------------------------------------- + +execGrapher :: HostOnlyFuns -> Grapher a -> (Graph, Sources, Sinks) +execGrapher hof m = + let s = R.runReader (execStateT m st) env + in (stateGraph s, stateSources s, stateSinks s) + where + env = + Env + { envHostOnlyFuns = hof, + envMeta = + Meta + { metaForkDepth = 0, + metaBodyDepth = 0, + metaGraphId = Nothing + } + } + st = + State + { stateGraph = MG.empty, + stateGraphedScalars = IS.empty, + stateSources = ([], []), + stateSinks = [], + stateUpdateAccs = IM.empty, + stateCopyableMemory = IM.empty, + stateStats = mempty + } + +-- | Execute a computation in a modified environment. +local :: (Env -> Env) -> Grapher a -> Grapher a +local f = mapStateT (R.local f) + +-- | Fetch the value of the environment. +ask :: Grapher Env +ask = lift R.ask + +-- | Retrieve a function of the current environment. +asks :: (Env -> a) -> Grapher a +asks = lift . R.asks + +-- | Register that the body contains a host-only statement. This means its +-- parent statement and any parent bodies themselves are host-only. A host-only +-- statement should not be migrated, either because it cannot run on device or +-- because it would be inefficient to do so. +tellHostOnly :: Grapher () +tellHostOnly = + modify $ \st -> st {stateStats = (stateStats st) {bodyHostOnly = True}} + +-- | Register that the body contains a GPUBody kernel. +tellGPUBody :: Grapher () +tellGPUBody = + modify $ \st -> st {stateStats = (stateStats st) {bodyHasGPUBody = True}} + +-- | Register that the current body contains a statement that reads device +-- memory. +tellRead :: Grapher () +tellRead = + modify $ \st -> st {stateStats = (stateStats st) {bodyReads = True}} + +-- | Register that these variables are used as operands within the current body. +tellOperands :: IdSet -> Grapher () +tellOperands is = + modify $ \st -> + let stats = stateStats st + operands = bodyOperands stats + in st {stateStats = stats {bodyOperands = operands <> is}} + +-- | Register that the current statement with a body at the given body depth is +-- host-only. +tellHostOnlyParent :: BodyDepth -> Grapher () +tellHostOnlyParent body_depth = + modify $ \st -> + let stats = stateStats st + parents = bodyHostOnlyParents stats + parents' = IS.insert body_depth parents + in st {stateStats = stats {bodyHostOnlyParents = parents'}} + +-- | Get the graph under construction. +getGraph :: Grapher Graph +getGraph = gets stateGraph + +-- | All scalar variables with a vertex representation in the graph. +getGraphedScalars :: Grapher IdSet +getGraphedScalars = gets stateGraphedScalars + +-- | Every known array that is backed by a memory segment that may be copied, +-- mapped to the outermost known body depth where an array is backed by a +-- superset of that segment. +-- +-- A body where all returned arrays are backed by such memory and are written by +-- its own statements will retain its asymptotic cost if migrated as a whole. +getCopyableMemory :: Grapher CopyableMemoryMap +getCopyableMemory = gets stateCopyableMemory + +-- | The outermost known body depth for an array backed by the same copyable +-- memory as the array with this name. +outermostCopyableArray :: VName -> Grapher (Maybe BodyDepth) +outermostCopyableArray n = IM.lookup (nameToId n) <$> getCopyableMemory + +-- | Reduces the variables to just the 'Id's of those that are scalars and which +-- have a vertex representation in the graph, excluding those that have been +-- connected to sinks. +onlyGraphedScalars :: Foldable t => t VName -> Grapher IdSet +onlyGraphedScalars vs = do + let is = foldl' (\s n -> IS.insert (nameToId n) s) IS.empty vs + IS.intersection is <$> getGraphedScalars + +-- | Like 'onlyGraphedScalars' but for a single 'VName'. +onlyGraphedScalar :: VName -> Grapher IdSet +onlyGraphedScalar n = do + let i = nameToId n + gss <- getGraphedScalars + if IS.member i gss + then pure (IS.singleton i) + else pure IS.empty + +-- | Like 'onlyGraphedScalars' but for a single t'SubExp'. +onlyGraphedScalarSubExp :: SubExp -> Grapher IdSet +onlyGraphedScalarSubExp (Constant _) = pure IS.empty +onlyGraphedScalarSubExp (Var n) = onlyGraphedScalar n + +-- | Update the graph under construction. +modifyGraph :: (Graph -> Graph) -> Grapher () +modifyGraph f = + modify $ \st -> st {stateGraph = f (stateGraph st)} + +-- | Update the contents of the graphed scalar set. +modifyGraphedScalars :: (IdSet -> IdSet) -> Grapher () +modifyGraphedScalars f = + modify $ \st -> st {stateGraphedScalars = f (stateGraphedScalars st)} + +-- | Update the contents of the copyable memory map. +modifyCopyableMemory :: (CopyableMemoryMap -> CopyableMemoryMap) -> Grapher () +modifyCopyableMemory f = + modify $ \st -> st {stateCopyableMemory = f (stateCopyableMemory st)} + +-- | Update the set of source connected vertices. +modifySources :: (Sources -> Sources) -> Grapher () +modifySources f = + modify $ \st -> st {stateSources = f (stateSources st)} + +-- | Record that this variable binds an array that is backed by copyable +-- memory shared by an array at this outermost body depth. +recordCopyableMemory :: Id -> BodyDepth -> Grapher () +recordCopyableMemory i bd = + modifyCopyableMemory (IM.insert i bd) + +-- | Increment the fork depth for variables graphed by this action. +incForkDepthFor :: Grapher a -> Grapher a +incForkDepthFor = + local $ \env -> + let meta = envMeta env + fork_depth = metaForkDepth meta + in env {envMeta = meta {metaForkDepth = fork_depth + 1}} + +-- | Increment the body depth for variables graphed by this action. +incBodyDepthFor :: Grapher a -> Grapher a +incBodyDepthFor = + local $ \env -> + let meta = envMeta env + body_depth = metaBodyDepth meta + in env {envMeta = meta {metaBodyDepth = body_depth + 1}} + +-- | Change the graph id for variables graphed by this action. +graphIdFor :: Id -> Grapher a -> Grapher a +graphIdFor i = + local $ \env -> + let meta = envMeta env + in env {envMeta = meta {metaGraphId = Just i}} + +-- | Capture body stats produced by the given action. +captureBodyStats :: Grapher a -> Grapher BodyStats +captureBodyStats m = do + stats <- gets stateStats + modify $ \st -> st {stateStats = mempty} + + _ <- m + + stats' <- gets stateStats + modify $ \st -> st {stateStats = stats <> stats'} + + pure stats' + +-- | Can applications of this function be moved to device? +isHostOnlyFun :: Name -> Grapher Bool +isHostOnlyFun fn = asks $ S.member fn . envHostOnlyFuns + +-- | Get the 'Meta' corresponding to the current body. +getMeta :: Grapher Meta +getMeta = asks envMeta + +-- | Get the body depth of the current body (its nesting level). +getBodyDepth :: Grapher BodyDepth +getBodyDepth = asks (metaBodyDepth . envMeta) diff --git a/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable/Graph.hs b/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable/Graph.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/ReduceDeviceSyncs/MigrationTable/Graph.hs @@ -0,0 +1,569 @@ +-- | This module contains the type definitions and basic operations +-- for the graph that +-- "Futhark.Optimise.ReduceDeviceSyncs.MigrationTable" internally uses +-- to construct a migration table. It is however completely +-- Futhark-agnostic and implements a generic graph abstraction. +-- +-- = Overview +-- +-- The 'Graph' type is a data flow dependency graph of program variables, each +-- variable represented by a 'Vertex'. A vertex may have edges to other vertices +-- or to a sink, which is a special vertex with no graph representation. Each +-- edge to a vertex is either from another vertex or from a source, which also +-- is a special vertex with no graph representation. +-- +-- The primary graph operation provided by this module is 'route'. Given the +-- vertex that some unspecified source has an edge to, a path is attempted +-- found to a sink. If a sink can be reached from the source, all edges along +-- the path are reversed. The path in the opposite direction of reversed edges +-- from a source to some sink is a route. +-- +-- Routes can be used to find a minimum vertex cut in the graph through what +-- amounts to a specialized depth-first search implementation of the +-- Ford-Fulkerson method. When viewed in this way each graph edge has a capacity +-- of 1 and the reversing of edges to create routes amounts to sending reverse +-- flow through a residual network (the current state of the graph). +-- The max-flow min-cut theorem allows one to determine a minimum edge cut that +-- separates the sources and sinks. +-- +-- If each vertex @v@ in the graph is viewed as two vertices, @v_in@ and +-- @v_out@, with all ingoing edges to @v@ going to @v_in@, all outgoing edges +-- from @v@ going from @v_out@, and @v_in@ connected to @v_out@ with a single +-- edge, then the minimum edge cut of the view amounts to a minimum vertex cut +-- in the actual graph. The view need not be manifested as whether @v_in@ or +-- @v_out@ is reached by an edge to @v@ can be determined from whether that edge +-- is reversed or not. The presence of an outgoing, reversed edge also gives the +-- state of the virtual edge that connects @v_in@ to @v_out@. +-- +-- When routing fails to find a sink in some subgraph reached via an edge then +-- that edge is marked exhausted. No sink can be reached via an exhausted edge, +-- and any subsequent routing attempt will skip pathfinding along such edge. +module Futhark.Optimise.ReduceDeviceSyncs.MigrationTable.Graph + ( -- * Types + Graph, + Id, + IdSet, + Vertex (..), + Routing (..), + Exhaustion (..), + Edges (..), + EdgeType (..), + Visited, + Result (..), + + -- * Construction + empty, + vertex, + declareEdges, + oneEdge, + none, + + -- * Insertion + insert, + + -- * Update + adjust, + connectToSink, + addEdges, + + -- * Query + member, + lookup, + isSinkConnected, + + -- * Routing + route, + routeMany, + + -- * Traversal + fold, + reduce, + ) +where + +import qualified Data.IntMap.Strict as IM +import qualified Data.IntSet as IS +import Data.List (foldl') +import qualified Data.Map.Strict as M +import Data.Maybe (fromJust) +import Prelude hiding (lookup) + +-------------------------------------------------------------------------------- +-- TYPES -- +-------------------------------------------------------------------------------- + +-- | A data flow dependency graph of program variables, each variable +-- represented by a 'Vertex'. +newtype Graph m = Graph (IM.IntMap (Vertex m)) + +-- | A handle that identifies a specific 'Vertex'. +type Id = Int + +-- | A set of 'Id's. +type IdSet = IS.IntSet + +-- | A graph representation of some program variable. +data Vertex m = Vertex + { -- | The handle for this vertex in the graph. + vertexId :: Id, + -- | Custom data associated with the variable. + vertexMeta :: m, + -- | Whether a route passes through this vertex, and from where. + vertexRouting :: Routing, + -- | Handles of vertices that this vertex has an edge to. + vertexEdges :: Edges + } + +-- | Route tracking for some vertex. +-- If a route passes through the vertex then both an ingoing and an outgoing +-- edge to/from that vertex will have been reversed, and the vertex will in +-- effect have lost one edge and gained another. The gained edge will be to +-- the prior vertex along the route that passes through. +data Routing + = -- | No route passes through the vertex, and no edges have been reversed, + -- added, nor deleted compared to what was declared. + NoRoute + | -- | A route passes through the vertex, and the prior vertex is the source + -- of that route. The edge gained by reversal is by definition exhausted. + FromSource + | -- | A route passes through the vertex, and this is the handle of the prior + -- vertex. The edge gained by reversal may be exhausted. Routing assumes + -- that at most one 'FromNode' routing exists to each vertex in a graph. + FromNode Id Exhaustion + deriving (Show, Eq, Ord) + +-- | Whether some edge is exhausted or not. No sink can be reached via an +-- exhausted edge. +data Exhaustion = Exhausted | NotExhausted + deriving (Show, Eq, Ord) + +-- | All relevant edges that have been declared from some vertex, plus +-- bookkeeping to track their exhaustion and reversal. +data Edges + = -- | The vertex has an edge to a sink; all other declared edges are + -- irrelevant. The edge cannot become exhausted, and it is reversed if a + -- route passes through the vertex (@vertexRouting v /= NoRoute@). + ToSink + | -- | All vertices that the vertex has a declared edge to, and which of + -- those edges that are not exhausted nor reversed, if not all. + ToNodes IdSet (Maybe IdSet) + deriving (Show, Eq, Ord) + +instance Semigroup Edges where + ToSink <> _ = ToSink + _ <> ToSink = ToSink + (ToNodes a1 Nothing) <> (ToNodes a2 Nothing) = + ToNodes (a1 <> a2) Nothing + (ToNodes a1 (Just e1)) <> (ToNodes a2 Nothing) = + ToNodes (a1 <> a2) $ Just (e1 <> IS.difference a2 a1) + (ToNodes a1 Nothing) <> (ToNodes a2 (Just e2)) = + ToNodes (a1 <> a2) $ Just (e2 <> IS.difference a1 a2) + (ToNodes a1 (Just e1)) <> (ToNodes a2 (Just e2)) = + let a = IS.difference e2 (IS.difference a1 e1) + b = IS.difference e1 (IS.difference a2 e2) + in ToNodes (a1 <> a2) $ Just (a <> b) + +instance Monoid Edges where + -- The empty set of edges. + mempty = ToNodes IS.empty Nothing + +-- | Whether a vertex is reached via a normal or reversed edge. +data EdgeType = Normal | Reversed + deriving (Eq, Ord) + +-- | State that tracks which vertices a traversal has visited, caching immediate +-- computations. +newtype Visited a = Visited {visited :: M.Map (EdgeType, Id) a} + +-- | The result of a graph traversal that may abort early in case a sink is +-- reached. +data Result a + = -- | The traversal finished without encountering a sink, producing this + -- value. + Produced a + | -- | The traversal was aborted because a sink was reached. + FoundSink + deriving (Eq) + +instance Semigroup a => Semigroup (Result a) where + FoundSink <> _ = FoundSink + _ <> FoundSink = FoundSink + Produced x <> Produced y = Produced (x <> y) + +-------------------------------------------------------------------------------- +-- CONSTRUCTION -- +-------------------------------------------------------------------------------- + +-- | The empty graph. +empty :: Graph m +empty = Graph IM.empty + +-- | Constructs a 'Vertex' without any edges. +vertex :: Id -> m -> Vertex m +vertex i m = + Vertex + { vertexId = i, + vertexMeta = m, + vertexRouting = NoRoute, + vertexEdges = mempty + } + +-- | Creates a set of edges where no edge is reversed or exhausted. +declareEdges :: [Id] -> Edges +declareEdges is = ToNodes (IS.fromList is) Nothing + +-- | Like 'declareEdges' but for a single vertex. +oneEdge :: Id -> Edges +oneEdge i = ToNodes (IS.singleton i) Nothing + +-- | Initial 'Visited' state before any vertex has been visited. +none :: Visited a +none = Visited M.empty + +-------------------------------------------------------------------------------- +-- INSERTION -- +-------------------------------------------------------------------------------- + +-- | Insert a new vertex into the graph. If its variable already is represented +-- in the graph, the original graph is returned. +insert :: Vertex m -> Graph m -> Graph m +insert v (Graph m) = Graph $ IM.insertWith const (vertexId v) v m + +-------------------------------------------------------------------------------- +-- UPDATE -- +-------------------------------------------------------------------------------- + +-- | Adjust the vertex with this specific id. When no vertex with that id is a +-- member of the graph, the original graph is returned. +adjust :: (Vertex m -> Vertex m) -> Id -> Graph m -> Graph m +adjust f i (Graph m) = Graph $ IM.adjust f i m + +-- | Connect the vertex with this id to a sink. When no vertex with that id is a +-- member of the graph, the original graph is returned. +connectToSink :: Id -> Graph m -> Graph m +connectToSink = adjust $ \v -> v {vertexEdges = ToSink} + +-- | Add these edges to the vertex with this id. When no vertex with that id is +-- a member of the graph, the original graph is returned. +addEdges :: Edges -> Id -> Graph m -> Graph m +addEdges es = adjust $ \v -> v {vertexEdges = es <> vertexEdges v} + +-------------------------------------------------------------------------------- +-- QUERY -- +-------------------------------------------------------------------------------- + +-- | Does a vertex for the given id exist in the graph? +member :: Id -> Graph m -> Bool +member i (Graph m) = IM.member i m + +-- | Returns the vertex with the given id. +lookup :: Id -> Graph m -> Maybe (Vertex m) +lookup i (Graph m) = IM.lookup i m + +-- | Returns whether a vertex with the given id exists in the +-- graph and is connected directly to a sink. +isSinkConnected :: Id -> Graph m -> Bool +isSinkConnected i g = + maybe False ((ToSink ==) . vertexEdges) (lookup i g) + +-- ROUTING -- +-------------------------------------------------------------------------------- + +-- | @route src g@ attempts to find a path in @g@ from the source connected +-- vertex with id @src@. If a sink is found, all edges along the path will be +-- reversed to create a route, and the id of the vertex that connects to the +-- sink is returned. +route :: Id -> Graph m -> (Maybe Id, Graph m) +route src g = + case route' IM.empty 0 Nothing Normal src g of + (DeadEnd, g') -> (Nothing, g') + (SinkFound snk, g') -> (Just snk, g') + (CycleDetected {}, _) -> + error + "Routing did not escape cycle in Futhark.Analysis.MigrationTable.Graph." + +-- | @routeMany srcs g@ attempts to create a 'route' in @g@ from every vertex +-- in @srcs@. Returns the ids for the vertices connected to each found sink. +routeMany :: [Id] -> Graph m -> ([Id], Graph m) +routeMany srcs graph = + foldl' f ([], graph) srcs + where + f (snks, g) src = + case route src g of + (Nothing, g') -> (snks, g') + (Just snk, g') -> (snk : snks, g') + +-------------------------------------------------------------------------------- +-- TRAVERSAL -- +-------------------------------------------------------------------------------- + +-- | @fold g f (a, vs) et i@ folds @f@ over the vertices in @g@ that can be +-- reached from the vertex with handle @i@ accessed via an edge of type @et@. +-- Each vertex @v@ may be visited up to two times, once for each type of edge +-- @e@ pointing to it, and each time @f a e v@ is evaluated to produce an +-- updated @a@ value to be used in future @f@ evaluations or to be returned. +-- The @vs@ set records which @f a e v@ evaluations already have taken place. +-- The function returns an updated 'Visited' set recording the evaluations it +-- has performed. +fold :: + Graph m -> + (a -> EdgeType -> Vertex m -> a) -> + (a, Visited ()) -> + EdgeType -> + Id -> + (a, Visited ()) +fold g f (res, vs) et i + | M.notMember (et, i) (visited vs), + Just v <- lookup i g = + let res' = f res et v + vs' = Visited $ M.insert (et, i) () (visited vs) + st = (res', vs') + in case (et, vertexRouting v) of + (Normal, FromSource) -> st + (Normal, FromNode rev _) -> foldReversed st rev + (Reversed, FromNode rev _) -> foldAll st rev (vertexEdges v) + _ -> foldNormals st (vertexEdges v) + | otherwise = + (res, vs) + where + foldReversed st = fold g f st Reversed + + foldAll st rev es = foldReversed (foldNormals st es) rev + + foldNormals st ToSink = st + foldNormals st (ToNodes es _) = + IS.foldl' (\s -> fold g f s Normal) st es + +-- | @reduce g r vs et i@ returns 'FoundSink' if a sink can be reached via the +-- vertex @v@ with id @i@ in @g@. Otherwise it returns 'Produced' @(r x et v)@ +-- where @x@ is the '<>' aggregate of all values produced by reducing the +-- vertices that are available via the edges of @v@. +-- @et@ identifies the type of edge that @v@ is accessed by and thereby which +-- edges of @v@ that are available. @vs@ caches reductions of vertices that +-- previously have been visited in the graph. +-- +-- The reduction of a cyclic reference resolves to 'mempty'. +reduce :: + Monoid a => + Graph m -> + (a -> EdgeType -> Vertex m -> a) -> + Visited (Result a) -> + EdgeType -> + Id -> + (Result a, Visited (Result a)) +reduce g r vs et i + | Just res <- M.lookup (et, i) (visited vs) = + (res, vs) + | Just v <- lookup i g = + reduceVertex v + | otherwise = + (Produced mempty, vs) -- shouldn't happen + where + reduceVertex v = + let (res, vs') = reduceEdges v + in case res of + Produced x -> cached (Produced $ r x et v) vs' + FoundSink -> cached res vs' + + cached res vs0 = + let vs1 = Visited (M.insert (et, i) res $ visited vs0) + in (res, vs1) + + reduceEdges v = + case (et, vertexRouting v) of + (Normal, FromSource) -> (Produced mempty, vs) + (Normal, FromNode rev _) -> entry (reduceReversed rev) + (Reversed, FromNode rev _) -> entry (reduceAll rev $ vertexEdges v) + _ -> entry (reduceNormals $ vertexEdges v) + + -- Handle cycles + entry f = f $ Visited $ M.insert (et, i) (Produced mempty) (visited vs) + + reduceReversed rev vs' = reduce g r vs' Reversed rev + + reduceAll rev es vs0 = + let (res, vs1) = reduceNormals es vs0 + in case res of + Produced _ -> + let (res', vs2) = reduceReversed rev vs1 + in (res <> res', vs2) + FoundSink -> (res, vs1) + + reduceNormals ToSink vs' = (FoundSink, vs') + reduceNormals (ToNodes es _) vs' = reduceNorms mempty (IS.elems es) vs' + + reduceNorms x [] vs0 = (Produced x, vs0) + reduceNorms x (e : es) vs0 = + let (res, vs1) = reduce g r vs0 Normal e + in case res of + Produced y -> reduceNorms (x <> y) es vs1 + FoundSink -> (res, vs1) + +-------------------------------------------------------------------------------- +-- ROUTING INTERNALS -- +-------------------------------------------------------------------------------- + +-- | A set of vertices visited by a graph traversal, and at what depth they were +-- encountered. Used to detect cycles. +type Pending = IM.IntMap Depth + +-- | Search depth. Distance to some vertex from some search root. +type Depth = Int + +-- | The outcome of attempted to find a route through a vertex. +data RoutingResult a + = -- | No sink could be reached through this vertex. + DeadEnd + | -- | A cycle was detected. A sink can be reached through this vertex if a + -- sink can be reached from the vertex at this depth. If no sink can be + -- reached from the vertex at this depth, then the graph should be updated + -- by these actions. Until the vertex is reached, the status of these + -- vertices are pending. + CycleDetected Depth [Graph a -> Graph a] Pending + | -- | A sink was found. This is the id of the vertex connected to it. + SinkFound Id + +instance Semigroup (RoutingResult a) where + SinkFound i <> _ = SinkFound i + _ <> SinkFound i = SinkFound i + CycleDetected d1 as1 _ <> CycleDetected d2 as2 p2 = + CycleDetected (min d1 d2) (as1 ++ as2) p2 + _ <> CycleDetected d as p = CycleDetected d as p + CycleDetected d as p <> _ = CycleDetected d as p + DeadEnd <> DeadEnd = DeadEnd + +instance Monoid (RoutingResult a) where + mempty = DeadEnd + +route' :: + Pending -> + Depth -> + Maybe Id -> + EdgeType -> + Id -> + Graph m -> + (RoutingResult m, Graph m) +route' p d prev et i g + | Just d' <- IM.lookup i p = + let found_cycle = (CycleDetected d' [] p, g) + in case et of + -- Accessing some vertex v via a normal edge corresponds to accessing + -- v_in via a normal edge. If v_in has a reversed edge then that is + -- the only outgoing edge that is available. + -- All outgoing edges available via this ingoing edge were already + -- available via the edge that first reached the vertex. + Normal -> found_cycle + -- Accessing some vertex v via a reversed edge corresponds to + -- accessing v_out via a reversed edge. All other edges of v_out are + -- available, and the edge from v_in to v_out has been reversed, + -- implying that v_in has a single reversed edge that also is + -- available. + -- There exists at most one reversed edge to each vertex. Since this + -- vertex was reached via one, and the vertex already have been + -- reached, then the first reach must have been via a normal edge + -- that only could traverse a reversed edge. The reversed edge from + -- v_out to v_in thus completes a cycle, but a sink might be + -- reachable via any of the other edges from v_out. + -- The depth for the vertex need not be updated as this is the only + -- edge to v_out and 'prev' is already in the 'Pending' map. + -- It follows that no (new) cycle can start here. + Reversed -> + let (res, g') = routeNormals (fromJust $ lookup i g) g p + in (fst found_cycle <> res, g') + | Just v <- lookup i g = + routeVertex v + | otherwise = + backtrack + where + backtrack = (DeadEnd, g) + + routeVertex v = + case (et, vertexRouting v) of + (Normal, FromSource) -> backtrack + (Normal, FromNode _ Exhausted) -> backtrack + (Normal, FromNode rev _) -> entry (routeReversed rev) + (Reversed, FromNode rev _) -> entry (routeAll rev v) + _ -> entry (routeNormals v) + + entry f = + let (res, g0) = f g (IM.insert i d p) + in case res of + CycleDetected d' as _ + | d == d' -> (DeadEnd, foldl' (\g1 a -> a g1) g0 as) + _ | otherwise -> (res, g0) + + routeAll rev v g0 p0 = + let (res, g1) = routeNormals v g0 p0 + in case res of + DeadEnd -> routeReversed rev g1 p0 + CycleDetected _ _ p1 -> + let (res', g2) = routeReversed rev g1 p1 + in (res <> res', g2) + SinkFound _ -> (res, g1) + + routeReversed rev g0 p0 = + let (res, g') = route' p0 (d + 1) (Just i) Reversed rev g0 + exhaust = flip adjust i $ + \v -> v {vertexRouting = FromNode rev Exhausted} + in case (res, et) of + (DeadEnd, _) -> + (res, exhaust g') + (CycleDetected d' as p', _) -> + (CycleDetected d' (exhaust : as) p', g') + (SinkFound _, Normal) -> + (res, setRoute g') + (SinkFound _, Reversed) -> + let f v = + v + { vertexEdges = withPrev (vertexEdges v), + vertexRouting = NoRoute + } + in (res, adjust f i g') + + setRoute = adjust (\v -> v {vertexRouting = routing}) i + + routing = + case prev of + Nothing -> FromSource + Just i' -> FromNode i' NotExhausted + + withPrev edges + | Just i' <- prev, + ToNodes es (Just es') <- edges = + ToNodes es (Just $ IS.insert i' es') + | otherwise = + edges -- shouldn't happen + routeNormals v g0 p0 + | ToSink <- vertexEdges v = + -- There cannot be a reversed edge to a vertex with an edge to a sink. + (SinkFound i, setRoute g0) + | ToNodes es nx <- vertexEdges v = + let (res, g', nx') = + case nx of + Just es' -> routeNorms (IS.toAscList es') g0 p0 + Nothing -> routeNorms (IS.toAscList es) g0 p0 + edges = ToNodes es (Just $ IS.fromDistinctAscList nx') + exhaust = flip adjust i $ \v' -> + v' {vertexEdges = ToNodes es (Just IS.empty)} + in case (res, et) of + (DeadEnd, _) -> (res, exhaust g') + (CycleDetected d' as p', _) -> + let res' = CycleDetected d' (exhaust : as) p' + v' = v {vertexEdges = edges} + in (res', insert v' g') + (SinkFound _, Normal) -> + let v' = v {vertexEdges = edges, vertexRouting = routing} + in (res, insert v' g') + (SinkFound _, Reversed) -> + let v' = v {vertexEdges = withPrev edges} + in (res, insert v' g') + + routeNorms [] g0 _ = (DeadEnd, g0, []) + routeNorms (e : es) g0 p0 = + let (res, g1) = route' p0 (d + 1) (Just i) Normal e g0 + in case res of + DeadEnd -> routeNorms es g1 p0 + SinkFound _ -> (res, g1, es) + CycleDetected _ _ p1 -> + let (res', g2, es') = routeNorms es g1 p1 + in (res <> res', g2, e : es') diff --git a/src/Futhark/Optimise/Simplify.hs b/src/Futhark/Optimise/Simplify.hs --- a/src/Futhark/Optimise/Simplify.hs +++ b/src/Futhark/Optimise/Simplify.hs @@ -45,7 +45,7 @@ -- We deepen the vtable so it will look like the constants are in an -- "outer loop"; this communicates useful information to some - -- simplification rules (e.g. seee issue #1302). + -- simplification rules (e.g. see issue #1302). funs' <- parPass (simplifyFun' (ST.deepen consts_vtable) . informFunDef) funs let funs_uses = UT.usages $ foldMap freeIn funs' diff --git a/src/Futhark/Optimise/Simplify/Engine.hs b/src/Futhark/Optimise/Simplify/Engine.hs --- a/src/Futhark/Optimise/Simplify/Engine.hs +++ b/src/Futhark/Optimise/Simplify/Engine.hs @@ -40,9 +40,11 @@ orIf, hasFree, isConsumed, + isConsuming, isFalse, isOp, isNotSafe, + isDeviceMigrated, asksEngineEnv, askVtable, localVtable, @@ -60,7 +62,9 @@ ST.SymbolTable, hoistStms, blockIf, + blockMigrated, enterLoop, + constructBody, module Futhark.Optimise.Simplify.Rep, ) where @@ -387,7 +391,7 @@ -- | Statements that are not worth hoisting out of loops, because they -- are unsafe, and added safety (by 'protectLoopHoisted') may inhibit --- further optimisation.. +-- further optimisation. notWorthHoisting :: ASTRep rep => BlockPred rep notWorthHoisting _ _ (Let pat _ e) = not (safeExp e) && any ((> 0) . arrayRank) (patTypes pat) @@ -402,7 +406,7 @@ cs' <- simplify cs (pat', pat_cs) <- collectCerts $ simplifyPat $ removePatWisdom pat let aux' = StmAux (cs' <> pat_cs) attrs dec - mkWiseLetStm pat' aux' <$> simplifyExpBase e + mkWiseStm pat' aux' <$> simplifyExpBase e -- Bottom-up simplify a statement. Recurses into sub-Bodies and Ops. -- Does not copy-propagate into the pattern and similar, as it is @@ -416,7 +420,7 @@ recSimplifyStm (Let pat (StmAux cs attrs (_, dec)) e) usage = do ((e', e_hoisted), e_cs) <- collectCerts $ simplifyExp usage pat e let aux' = StmAux (cs <> e_cs) attrs dec - pure (e_hoisted, mkWiseLetStm (removePatWisdom pat) aux' e') + pure (e_hoisted, mkWiseStm (removePatWisdom pat) aux' e') hoistStms :: SimplifiableRep rep => @@ -860,7 +864,7 @@ Nothing -> pure (Nothing, mempty) Just (op_lam, nes) -> do - (op_lam', op_lam_stms) <- simplifyLambda op_lam + (op_lam', op_lam_stms) <- blockMigrated (simplifyLambda op_lam) nes' <- simplify nes pure (Just (op_lam', nes'), op_lam_stms) (,op_stms) <$> ((,,op') <$> simplify shape <*> simplify arrs) @@ -870,6 +874,36 @@ simplifyExp _ _ e = do e' <- simplifyExpBase e pure (e', mempty) + +-- | Block hoisting of 'Index' statements introduced by migration. +blockMigrated :: + SimplifiableRep rep => + SimpleM rep (Lambda (Wise rep), Stms (Wise rep)) -> + SimpleM rep (Lambda (Wise rep), Stms (Wise rep)) +blockMigrated = local withMigrationBlocker + where + withMigrationBlocker (ops, env) = + let blockers = envHoistBlockers env + par_blocker = blockHoistPar blockers + + blocker = par_blocker `orIf` isDeviceMigrated + + blockers' = blockers {blockHoistPar = blocker} + env' = env {envHoistBlockers = blockers'} + in (ops, env') + +-- | Statement is a scalar read from a single element array of rank one. +isDeviceMigrated :: SimplifiableRep rep => BlockPred (Wise rep) +isDeviceMigrated vtable _ stm + | BasicOp (Index arr slice) <- stmExp stm, + [DimFix idx] <- unSlice slice, + idx == intConst Int64 0, + Just arr_t <- ST.lookupType arr vtable, + [size] <- arrayDims arr_t, + size == intConst Int64 1 = + True + | otherwise = + False -- The simple nonrecursive case that we can perform without bottom-up -- information. diff --git a/src/Futhark/Optimise/Simplify/Rep.hs b/src/Futhark/Optimise/Simplify/Rep.hs --- a/src/Futhark/Optimise/Simplify/Rep.hs +++ b/src/Futhark/Optimise/Simplify/Rep.hs @@ -22,7 +22,7 @@ addScopeWisdom, addWisdomToPat, mkWiseBody, - mkWiseLetStm, + mkWiseStm, mkWiseExpDec, CanBeWise (..), @@ -31,6 +31,7 @@ informLambda, informFunDef, informStms, + informBody, ) where @@ -240,13 +241,13 @@ (aliases, consumed) = Aliases.mkBodyAliasing stms res -- | Produce a statement with simplifier information. -mkWiseLetStm :: +mkWiseStm :: (ASTRep rep, CanBeWise (Op rep)) => Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp (Wise rep) -> Stm (Wise rep) -mkWiseLetStm pat (StmAux cs attrs dec) e = +mkWiseStm pat (StmAux cs attrs dec) e = let pat' = addWisdomToPat pat e in Let pat' (StmAux cs attrs $ mkWiseExpDec pat' dec e) e @@ -275,7 +276,7 @@ env <- asksScope removeScopeWisdom flip runReaderT env $ do Let pat dec _ <- mkLetNames names $ removeExpWisdom e - pure $ mkWiseLetStm pat dec e + pure $ mkWiseStm pat dec e mkBody stms res = let Body bodyrep _ _ = mkBody (fmap removeStmWisdom stms) res @@ -298,7 +299,7 @@ -- | Construct a 'Wise' statement. informStm :: Informing rep => Stm rep -> Stm (Wise rep) -informStm (Let pat aux e) = mkWiseLetStm pat aux $ informExp e +informStm (Let pat aux e) = mkWiseStm pat aux $ informExp e -- | Construct 'Wise' statements. informStms :: Informing rep => Stms rep -> Stms (Wise rep) diff --git a/src/Futhark/Optimise/Sink.hs b/src/Futhark/Optimise/Sink.hs --- a/src/Futhark/Optimise/Sink.hs +++ b/src/Futhark/Optimise/Sink.hs @@ -49,8 +49,12 @@ import Data.Bifunctor import Data.List (foldl') import qualified Data.Map as M +import Data.Sequence ((<|)) +import qualified Data.Sequence as SQ import qualified Futhark.Analysis.Alias as Alias import qualified Futhark.Analysis.SymbolTable as ST +import Futhark.Builder.Class +import Futhark.Construct (sliceDim) import Futhark.IR.Aliases import Futhark.IR.GPU import Futhark.IR.MC @@ -67,6 +71,7 @@ type Constraints rep = ( ASTRep rep, Aliased rep, + Buildable rep, ST.IndexOp (Op rep) ) @@ -103,6 +108,39 @@ && all (`ST.available` vtable) (namesToList (freeIn stm)) sunk = namesFromList $ foldMap (patNames . stmPat) sunk_stms +optimiseLoop :: + Constraints rep => + Sinker rep (Op rep) -> + Sinker rep ([(FParam rep, SubExp)], LoopForm rep, Body rep) +optimiseLoop onOp vtable sinking (merge, form, body0) + | WhileLoop {} <- form = + let (body1, sunk) = optimiseBody onOp vtable' sinking body0 + in ((merge, form, body1), sunk) + | ForLoop i it bound loop_vars <- form = + let stms' = foldr (inline i) (bodyStms body0) loop_vars + body1 = body0 {bodyStms = stms'} + (body2, sunk) = optimiseBody onOp vtable' sinking body1 + notSunk (x, _) = not $ paramName x `nameIn` sunk + loop_vars' = filter notSunk loop_vars + form' = ForLoop i it bound loop_vars' + body3 = body2 {bodyStms = SQ.drop (length loop_vars') (bodyStms body2)} + in ((merge, form', body3), sunk) + where + (params, _) = unzip merge + scope = case form of + WhileLoop {} -> scopeOfFParams params + ForLoop i it _ _ -> M.insert i (IndexName it) $ scopeOfFParams params + vtable' = ST.fromScope scope <> vtable + + inline i (x, arr) stms = + let pt = typeOf x + slice = Slice $ DimFix (Var i) : map sliceDim (arrayDims pt) + e = BasicOp (Index arr slice) + pat = mkExpPat [Ident (paramName x) pt] e + aux = StmAux mempty mempty (mkExpDec pat e) + stm = Let pat aux e + in stm <| stms + optimiseStms :: Constraints rep => Sinker rep (Op rep) -> @@ -140,6 +178,15 @@ in ( stm {stmExp = If cond tbranch' fbranch' ret} : stms', tsunk <> fsunk <> sunk ) + | DoLoop merge lform body <- stmExp stm = + let comps = (merge, lform, body) + (comps', loop_sunk) = optimiseLoop onOp vtable sinking comps + (merge', lform', body') = comps' + + (stms', stms_sunk) = optimiseStms' vtable' sinking stms + in ( stm {stmExp = DoLoop merge' lform' body'} : stms', + stms_sunk <> loop_sunk + ) | Op op <- stmExp stm = let (op', op_sunk) = onOp vtable sinking op (stms', stms_sunk) = optimiseStms' vtable' sinking stms @@ -211,7 +258,7 @@ type SinkRep rep = Aliases rep sink :: - ( ASTRep rep, + ( Buildable rep, CanBeAliased (Op rep), ST.IndexOp (OpWithAliases (Op rep)) ) => @@ -241,6 +288,8 @@ onHostOp :: Sinker (SinkRep GPU) (Op (SinkRep GPU)) onHostOp vtable sinking (SegOp op) = first SegOp $ optimiseSegOp onHostOp vtable sinking op + onHostOp vtable sinking (GPUBody types body) = + first (GPUBody types) $ optimiseBody onHostOp vtable sinking body onHostOp _ _ op = (op, mempty) -- | Sinking for multicore. diff --git a/src/Futhark/Optimise/TileLoops.hs b/src/Futhark/Optimise/TileLoops.hs --- a/src/Futhark/Optimise/TileLoops.hs +++ b/src/Futhark/Optimise/TileLoops.hs @@ -31,35 +31,45 @@ onStms scope stms = modifyNameSource $ runState $ - runReaderT (optimiseStms stms) scope + runReaderT (optimiseStms (M.empty, M.empty) stms) scope -optimiseBody :: Body GPU -> TileM (Body GPU) -optimiseBody (Body () stms res) = - Body () <$> optimiseStms stms <*> pure res +optimiseBody :: Env -> Body GPU -> TileM (Body GPU) +optimiseBody env (Body () stms res) = + Body () <$> optimiseStms env stms <*> pure res -optimiseStms :: Stms GPU -> TileM (Stms GPU) -optimiseStms stms = - localScope (scopeOf stms) $ - mconcat <$> mapM optimiseStm (stmsToList stms) +optimiseStms :: Env -> Stms GPU -> TileM (Stms GPU) +optimiseStms env stms = + localScope (scopeOf stms) $ do + (_, stms') <- foldM foldfun (env, mempty) $ stmsToList stms + pure stms' + where + foldfun :: (Env, Stms GPU) -> Stm GPU -> TileM (Env, Stms GPU) + foldfun (e, ss) s = do + (e', s') <- optimiseStm e s + pure (e', ss <> s') -optimiseStm :: Stm GPU -> TileM (Stms GPU) -optimiseStm stm@(Let pat aux (Op (SegOp (SegMap lvl@SegThread {} space ts kbody)))) = do - res3dtiling <- doRegTiling3D stm - case res3dtiling of - Just (extra_stms, stmt') -> pure (extra_stms <> oneStm stmt') - Nothing -> do - blkRegTiling_res <- mmBlkRegTiling stm - case blkRegTiling_res of - Just (extra_stms, stmt') -> pure (extra_stms <> oneStm stmt') - Nothing -> localScope (scopeOfSegSpace space) $ do - (host_stms, (lvl', space', kbody')) <- tileInKernelBody mempty initial_variance lvl space ts kbody - pure $ host_stms <> oneStm (Let pat aux $ Op $ SegOp $ SegMap lvl' space' ts kbody') +optimiseStm :: Env -> Stm GPU -> TileM (Env, Stms GPU) +optimiseStm env stm@(Let pat aux (Op (SegOp (SegMap lvl@SegThread {} space ts kbody)))) = do + res3dtiling <- localScope (scopeOfSegSpace space) $ doRegTiling3D stm + stms' <- + case res3dtiling of + Just (extra_stms, stmt') -> pure (extra_stms <> oneStm stmt') + Nothing -> do + blkRegTiling_res <- mmBlkRegTiling env stm + case blkRegTiling_res of + Just (extra_stms, stmt') -> pure (extra_stms <> oneStm stmt') + Nothing -> localScope (scopeOfSegSpace space) $ do + (host_stms, (lvl', space', kbody')) <- tileInKernelBody mempty initial_variance lvl space ts kbody + pure $ host_stms <> oneStm (Let pat aux $ Op $ SegOp $ SegMap lvl' space' ts kbody') + pure (env, stms') where initial_variance = M.map mempty $ scopeOfSegSpace space -optimiseStm (Let pat aux e) = - pure <$> (Let pat aux <$> mapExpM optimise e) +optimiseStm env (Let pat aux e) = do + env' <- changeEnv env (head $ patNames pat) e + e' <- mapExpM (optimise env') e + pure (env', oneStm $ Let pat aux e') where - optimise = identityMapper {mapOnBody = \scope -> localScope scope . optimiseBody} + optimise env' = identityMapper {mapOnBody = \scope -> localScope scope . optimiseBody env'} tileInKernelBody :: Names -> @@ -618,9 +628,6 @@ <$> (zip <$> mapM lookupType m_body_free <*> pure m_body_free) let blank t = maybe (eBlank t) (pure . BasicOp . SubExp . Var) $ lookup t t_to_v letTupExp desc =<< eIf (toExp in_bounds) (pure m_body) (eBody $ map blank ts) - where - isAcc Acc {} = True - isAcc _ = False postludeGeneric :: Tiling -> diff --git a/src/Futhark/Optimise/TileLoops/Shared.hs b/src/Futhark/Optimise/TileLoops/Shared.hs --- a/src/Futhark/Optimise/TileLoops/Shared.hs +++ b/src/Futhark/Optimise/TileLoops/Shared.hs @@ -1,5 +1,12 @@ +{-# LANGUAGE FlexibleContexts #-} + module Futhark.Optimise.TileLoops.Shared ( TileM, + Env, + index, + update, + forLoop', + forLoop, segMap1D, segMap2D, segMap3D, @@ -7,6 +14,7 @@ VarianceTable, varianceInStms, isTileableRedomap, + changeEnv, TileKind (..), ) where @@ -16,12 +24,64 @@ import Data.List (foldl', zip4) import qualified Data.Map as M import Futhark.IR.GPU +import qualified Futhark.IR.Mem.IxFun as IxFun +import qualified Futhark.IR.SeqMem as ExpMem import Futhark.MonadFreshNames import Futhark.Tools import Futhark.Transform.Rename type TileM = ReaderT (Scope GPU) (State VNameSource) +-- | Are we working with full or partial tiles? +data TileKind = TilePartial | TileFull + +-- index an array with indices given in outer_indices; any inner +-- dims of arr not indexed by outer_indices are sliced entirely +index :: MonadBuilder m => String -> VName -> [VName] -> m VName +index se_desc arr outer_indices = do + arr_t <- lookupType arr + let shape = arrayShape arr_t + inner_dims = shapeDims $ stripDims (length outer_indices) shape + untouched d = DimSlice (intConst Int64 0) d (intConst Int64 1) + inner_slices = map untouched inner_dims + slice = Slice $ map (DimFix . Var) outer_indices ++ inner_slices + letExp se_desc $ BasicOp $ Index arr slice + +update :: MonadBuilder m => String -> VName -> [VName] -> SubExp -> m VName +update se_desc arr indices new_elem = + letExp se_desc $ BasicOp $ Update Unsafe arr (Slice $ map (DimFix . Var) indices) new_elem + +forLoop' :: + SubExp -> -- loop var + [VName] -> -- loop inits + ( VName -> + [VName] -> -- (loop var -> loop inits -> loop body) + Builder GPU (Body GPU) + ) -> + Builder GPU [VName] +forLoop' i_bound merge body = do + i <- newVName "i" -- could give this as arg to the function + let loop_form = ForLoop i Int64 i_bound [] + + merge_ts <- mapM lookupType merge + loop_inits <- mapM (\merge_t -> newParam "merge" $ toDecl merge_t Unique) merge_ts + + loop_body <- + runBodyBuilder . inScopeOf loop_form . localScope (scopeOfFParams loop_inits) $ + body i $ map paramName loop_inits + + letTupExp "loop" $ + DoLoop (zip loop_inits $ map Var merge) loop_form loop_body + +forLoop :: + SubExp -> + [VName] -> + (VName -> [VName] -> Builder GPU (Body GPU)) -> + Builder GPU VName +forLoop i_bound merge body = do + res_list <- forLoop' i_bound merge body + pure $ head res_list + segMap1D :: String -> SegLevel -> @@ -41,8 +101,7 @@ let ret (SubExpRes cs se) = Returns manifest cs se letTupExp desc $ - Op . SegOp $ - SegMap lvl space ts $ KernelBody () stms' $ map ret res' + Op . SegOp $ SegMap lvl space ts $ KernelBody () stms' $ map ret res' segMap2D :: String -> -- desc @@ -55,8 +114,8 @@ Builder GPU [VName] segMap2D desc lvl manifest (dim_y, dim_x) f = do ltid_xx <- newVName "ltid_x" - ltid_flat <- newVName "ltid_flat" ltid_yy <- newVName "ltid_y" + ltid_flat <- newVName "ltid_flat" let segspace = SegSpace ltid_flat [(ltid_yy, dim_y), (ltid_xx, dim_x)] ((ts, res), stms) <- localScope (scopeOfSegSpace segspace) . runBuilder $ do @@ -66,8 +125,7 @@ let ret (SubExpRes cs se) = Returns manifest cs se letTupExp desc <=< renameExp $ - Op . SegOp $ - SegMap lvl segspace ts $ KernelBody () stms $ map ret res + Op . SegOp $ SegMap lvl segspace ts $ KernelBody () stms $ map ret res segMap3D :: String -> -- desc @@ -79,10 +137,10 @@ ) -> Builder GPU [VName] segMap3D desc lvl manifest (dim_z, dim_y, dim_x) f = do - ltid_x <- newVName "ltid_x" ltid_flat <- newVName "ltid_flat" - ltid_y <- newVName "ltid_y" ltid_z <- newVName "ltid_z" + ltid_y <- newVName "ltid_y" + ltid_x <- newVName "ltid_x" let segspace = SegSpace ltid_flat [(ltid_z, dim_z), (ltid_y, dim_y), (ltid_x, dim_x)] ((ts, res), stms) <- localScope (scopeOfSegSpace segspace) . runBuilder $ do @@ -92,12 +150,11 @@ let ret (SubExpRes cs se) = Returns manifest cs se letTupExp desc <=< renameExp $ - Op . SegOp $ - SegMap lvl segspace ts $ KernelBody () stms $ map ret res + Op . SegOp $ SegMap lvl segspace ts $ KernelBody () stms $ map ret res segScatter2D :: - String -> -- desc - SubExp -> -- arr_size + String -> + SubExp -> VName -> SegLevel -> -- lvl [SubExp] -> -- dims of sequential loop on top @@ -105,14 +162,14 @@ ([VName] -> (VName, VName) -> Builder GPU (SubExp, SubExp)) -> -- f Builder GPU VName segScatter2D desc arr_size updt_arr lvl seq_dims (dim_x, dim_y) f = do - ltid_x <- newVName "ltid_x" - ltid_y <- newVName "ltid_y" ltid_flat <- newVName "ltid_flat" + ltid_y <- newVName "ltid_y" + ltid_x <- newVName "ltid_x" seq_is <- replicateM (length seq_dims) (newVName "ltid_seq") let seq_space = zip seq_is seq_dims - let segspace = SegSpace ltid_flat $ seq_space ++ [(ltid_x, dim_x), (ltid_y, dim_y)] + let segspace = SegSpace ltid_flat $ seq_space ++ [(ltid_y, dim_y), (ltid_x, dim_x)] lvl' = SegThread (segNumGroups lvl) @@ -122,7 +179,7 @@ ((t_v, res_v, res_i), stms) <- runBuilder $ do (res_v, res_i) <- localScope (scopeOfSegSpace segspace) $ - f seq_is (ltid_x, ltid_y) + f seq_is (ltid_y, ltid_x) t_v <- subExpType res_v pure (t_v, res_v, res_i) @@ -195,5 +252,71 @@ varianceInStms :: VarianceTable -> Stms GPU -> VarianceTable varianceInStms = foldl' varianceInStm --- | Are we working with full or partial tiles? -data TileKind = TilePartial | TileFull +---------------- +---- Helpers for building the environment that binds array variable names to their index functions +---------------- + +type IxFun = IxFun.IxFun (TPrimExp Int64 VName) + +-- | Map from array variable names to their corresponding index functions. +-- The info is not guaranteed to be exact, e.g., we assume ifs and loops +-- return arrays layed out in normalized (row-major) form in memory. +-- We only record aliasing statements, such as transposition, slice, etc. +type IxFnEnv = M.Map VName IxFun + +type WithEnv = M.Map VName (Lambda GPU, [SubExp]) + +type Env = (WithEnv, IxFnEnv) + +changeEnv :: Env -> VName -> Exp GPU -> TileM Env +changeEnv (with_env, ixfn_env) y e = do + with_env' <- changeWithEnv with_env e + ixfn_env' <- changeIxFnEnv ixfn_env y e + pure (with_env', ixfn_env') + +changeWithEnv :: WithEnv -> Exp GPU -> TileM WithEnv +changeWithEnv with_env (WithAcc accum_decs inner_lam) = do + let bindings = map mapfun accum_decs + par_tps = take (length bindings) $ map paramName $ lambdaParams inner_lam + with_env' = M.union with_env $ M.fromList $ zip par_tps bindings + pure with_env' + where + mapfun (_, _, Nothing) = error "What the hack is an accumulator without operator?" + mapfun (shp, _, Just (lam_inds, ne)) = + let len_inds = length $ shapeDims shp + lam_op = lam_inds {lambdaParams = drop len_inds $ lambdaParams lam_inds} + in (lam_op, ne) +changeWithEnv with_env _ = pure with_env + +composeIxfuns :: IxFnEnv -> VName -> VName -> (IxFun -> IxFun) -> TileM IxFnEnv +composeIxfuns env y x ixf_fun = + case M.lookup x env of + Just ixf -> pure $ M.insert y (ixf_fun ixf) env + Nothing -> do + tp <- lookupType x + case tp of + Array _ptp shp _u -> do + let shp' = map ExpMem.pe64 (shapeDims shp) + pure $ M.insert y (ixf_fun $ IxFun.iota shp') env + _ -> pure env + +changeIxFnEnv :: IxFnEnv -> VName -> Exp GPU -> TileM IxFnEnv +changeIxFnEnv env y (BasicOp (Reshape shp_chg x)) = + composeIxfuns env y x (`IxFun.reshape` map (fmap ExpMem.pe64) shp_chg) +changeIxFnEnv env y (BasicOp (Manifest perm x)) = do + tp <- lookupType x + case tp of + Array _ptp shp _u -> do + let shp' = map ExpMem.pe64 (shapeDims shp) + let ixfn = IxFun.permute (IxFun.iota shp') perm + pure $ M.insert y ixfn env + _ -> error "In TileLoops/Shared.hs, changeIxFnEnv: manifest applied to a non-array!" +changeIxFnEnv env y (BasicOp (Rearrange perm x)) = + composeIxfuns env y x (`IxFun.permute` perm) +changeIxFnEnv env y (BasicOp (Rotate rs x)) = + composeIxfuns env y x (`IxFun.rotate` fmap ExpMem.pe64 rs) +changeIxFnEnv env y (BasicOp (Index x slc)) = + composeIxfuns env y x (`IxFun.slice` (Slice $ map (fmap ExpMem.pe64) $ unSlice slc)) +changeIxFnEnv env y (BasicOp (Opaque _ (Var x))) = + composeIxfuns env y x id +changeIxFnEnv env _ _ = pure env diff --git a/src/Futhark/Pass.hs b/src/Futhark/Pass.hs --- a/src/Futhark/Pass.hs +++ b/src/Futhark/Pass.hs @@ -95,9 +95,9 @@ let ((x', log), src') = runState (runPassM (f a)) src in (x', log, src') --- | Apply some operation to the top-level constants. Then applies an --- operation to all the function function definitions, which are also --- given the transformed constants so they can be brought into scope. +-- | Apply some operation to the top-level constants. Then applies an +-- operation to all the function definitions, which are also given the +-- transformed constants so they can be brought into scope. -- The function definition transformations are run in parallel (with -- 'parPass'), since they cannot affect each other. intraproceduralTransformationWithConsts :: diff --git a/src/Futhark/Pass/AD.hs b/src/Futhark/Pass/AD.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Pass/AD.hs @@ -0,0 +1,102 @@ +{-# LANGUAGE TypeFamilies #-} + +-- | Apply all AD operators in the program, leaving AD-free code. +module Futhark.Pass.AD (applyAD, applyADInnermost) where + +import Control.Monad +import Control.Monad.Reader +import Futhark.AD.Fwd (fwdJVP) +import Futhark.AD.Rev (revVJP) +import Futhark.Builder +import Futhark.IR.SOACS +import Futhark.IR.SOACS.Simplify (simplifyLambda) +import Futhark.Pass + +-- | Whether we apply only the innermost AD operators, or all of them. +-- The former is very useful for debugging, but probably not useful +-- for actual compilation. +data Mode = Innermost | All + deriving (Eq) + +bindLambda :: + (MonadBuilder m, Rep m ~ SOACS) => + Pat Type -> + StmAux (ExpDec SOACS) -> + Lambda SOACS -> + [SubExp] -> + m () +bindLambda pat aux (Lambda params body _) args = do + auxing aux . forM_ (zip params args) $ \(param, arg) -> + letBindNames [paramName param] $ + BasicOp $ case (paramType param, arg) of + (Array {}, Var v) -> Copy v + _ -> SubExp arg + res <- bodyBind body + forM_ (zip (patNames pat) res) $ \(v, SubExpRes cs se) -> + certifying cs $ letBindNames [v] $ BasicOp $ SubExp se + +onStm :: Mode -> Scope SOACS -> Stm SOACS -> PassM (Stms SOACS) +onStm mode scope (Let pat aux (Op (VJP lam args vec))) = do + lam' <- onLambda mode scope lam + if mode == All || lam == lam' + then do + lam'' <- (`runReaderT` scope) . simplifyLambda =<< revVJP scope lam' + runBuilderT_ (bindLambda pat aux lam'' $ args ++ vec) scope + else pure $ oneStm $ Let pat aux $ Op $ VJP lam' args vec +onStm mode scope (Let pat aux (Op (JVP lam args vec))) = do + lam' <- onLambda mode scope lam + if mode == All || lam == lam' + then do + lam'' <- fwdJVP scope lam' + runBuilderT_ (bindLambda pat aux lam'' $ args ++ vec) scope + else pure $ oneStm $ Let pat aux $ Op $ JVP lam' args vec +onStm mode scope (Let pat aux e) = oneStm . Let pat aux <$> mapExpM mapper e + where + mapper = + identityMapper + { mapOnBody = \bscope -> onBody mode (bscope <> scope), + mapOnOp = mapSOACM soac_mapper + } + soac_mapper = identitySOACMapper {mapOnSOACLambda = onLambda mode scope} + +onStms :: Mode -> Scope SOACS -> Stms SOACS -> PassM (Stms SOACS) +onStms mode scope stms = mconcat <$> mapM (onStm mode scope') (stmsToList stms) + where + scope' = scopeOf stms <> scope + +onBody :: Mode -> Scope SOACS -> Body SOACS -> PassM (Body SOACS) +onBody mode scope body = do + stms <- onStms mode scope $ bodyStms body + pure $ body {bodyStms = stms} + +onLambda :: Mode -> Scope SOACS -> Lambda SOACS -> PassM (Lambda SOACS) +onLambda mode scope lam = do + body <- onBody mode (scopeOfLParams (lambdaParams lam) <> scope) $ lambdaBody lam + pure $ lam {lambdaBody = body} + +onFun :: Mode -> Stms SOACS -> FunDef SOACS -> PassM (FunDef SOACS) +onFun mode consts fd = do + body <- onBody mode (scopeOf consts <> scopeOf fd) $ funDefBody fd + pure $ fd {funDefBody = body} + +applyAD :: Pass SOACS SOACS +applyAD = + Pass + { passName = "ad", + passDescription = "Apply AD operators", + passFunction = + intraproceduralTransformationWithConsts + (onStms All mempty) + (onFun All) + } + +applyADInnermost :: Pass SOACS SOACS +applyADInnermost = + Pass + { passName = "ad innermost", + passDescription = "Apply innermost AD operators", + passFunction = + intraproceduralTransformationWithConsts + (onStms Innermost mempty) + (onFun Innermost) + } diff --git a/src/Futhark/Pass/ExpandAllocations.hs b/src/Futhark/Pass/ExpandAllocations.hs --- a/src/Futhark/Pass/ExpandAllocations.hs +++ b/src/Futhark/Pass/ExpandAllocations.hs @@ -735,6 +735,7 @@ unAllocOp Alloc {} = Left "unAllocOp: unhandled Alloc" unAllocOp (Inner OtherOp {}) = Left "unAllocOp: unhandled OtherOp" + unAllocOp (Inner GPUBody {}) = Left "unAllocOp: unhandled GPUBody" unAllocOp (Inner (SizeOp op)) = pure $ SizeOp op unAllocOp (Inner (SegOp op)) = SegOp <$> mapSegOpM mapper op where diff --git a/src/Futhark/Pass/ExplicitAllocations/GPU.hs b/src/Futhark/Pass/ExplicitAllocations/GPU.hs --- a/src/Futhark/Pass/ExplicitAllocations/GPU.hs +++ b/src/Futhark/Pass/ExplicitAllocations/GPU.hs @@ -76,6 +76,8 @@ error $ "Cannot allocate memory in SOAC: " ++ pretty op handleHostOp (SegOp op) = Inner . SegOp <$> handleSegOp op +handleHostOp (GPUBody ts (Body _ stms res)) = + fmap (Inner . GPUBody ts) . buildBody_ . allocInStms stms $ pure res kernelExpHints :: Exp GPUMem -> AllocM GPU GPUMem [ExpHint] kernelExpHints (BasicOp (Manifest perm v)) = do diff --git a/src/Futhark/Pass/ExtractMulticore.hs b/src/Futhark/Pass/ExtractMulticore.hs --- a/src/Futhark/Pass/ExtractMulticore.hs +++ b/src/Futhark/Pass/ExtractMulticore.hs @@ -284,6 +284,10 @@ pure (red_stms, op) transformSOAC :: Pat Type -> Attrs -> SOAC SOACS -> ExtractM (Stms MC) +transformSOAC _ _ JVP {} = + error "transformSOAC: unhandled JVP" +transformSOAC _ _ VJP {} = + error "transformSOAC: unhandled VJP" transformSOAC pat _ (Screma w arrs form) | Just lam <- isMapSOAC form = do seq_op <- transformMap DoNotRename sequentialiseBody w lam arrs diff --git a/src/Futhark/Passes.hs b/src/Futhark/Passes.hs --- a/src/Futhark/Passes.hs +++ b/src/Futhark/Passes.hs @@ -17,18 +17,23 @@ import Futhark.IR.GPUMem (GPUMem) import Futhark.IR.MC (MC) import Futhark.IR.MCMem (MCMem) -import Futhark.IR.SOACS (SOACS) +import Futhark.IR.SOACS (SOACS, usesAD) import Futhark.IR.Seq (Seq) import Futhark.IR.SeqMem (SeqMem) import Futhark.Optimise.CSE import Futhark.Optimise.DoubleBuffer import Futhark.Optimise.Fusion +import Futhark.Optimise.GenRedOpt +import Futhark.Optimise.HistAccs import Futhark.Optimise.InPlaceLowering import Futhark.Optimise.InliningDeadFun import qualified Futhark.Optimise.MemoryBlockMerging as MemoryBlockMerging +import Futhark.Optimise.MergeGPUBodies +import Futhark.Optimise.ReduceDeviceSyncs import Futhark.Optimise.Sink import Futhark.Optimise.TileLoops import Futhark.Optimise.Unstream +import Futhark.Pass.AD import Futhark.Pass.ExpandAllocations import qualified Futhark.Pass.ExplicitAllocations.GPU as GPU import qualified Futhark.Pass.ExplicitAllocations.MC as MC @@ -62,7 +67,23 @@ simplifySOACS, removeDeadFunctions ] + >>> condPipeline usesAD adPipeline +-- | This is the pipeline that applies the AD transformation and +-- subsequent interesting optimisations. +adPipeline :: Pipeline SOACS SOACS +adPipeline = + passes + [ applyAD, + simplifySOACS, + performCSE True, + fuseSOACs, + performCSE True, + simplifySOACS, + fuseSOACs, + simplifySOACS + ] + -- | The pipeline used by the CUDA and OpenCL backends, but before -- adding memory information. Includes 'standardPipeline'. kernelsPipeline :: Pipeline SOACS GPU @@ -71,12 +92,23 @@ >>> onePass extractKernels >>> passes [ simplifyGPU, - babysitKernels, + optimiseGenRed, + simplifyGPU, tileLoops, + simplifyGPU, + histAccsGPU, + babysitKernels, + simplifyGPU, unstreamGPU, performCSE True, simplifyGPU, - sinkGPU, + sinkGPU, -- Sink reads before migrating them. + reduceDeviceSyncs, + simplifyGPU, -- Simplify and hoist storages. + performCSE True, -- Eliminate duplicate storages. + mergeGPUBodies, + simplifyGPU, -- Cleanup merged GPUBody kernels. + sinkGPU, -- Sink reads within GPUBody kernels. inPlaceLoweringGPU ] diff --git a/src/Futhark/Pipeline.hs b/src/Futhark/Pipeline.hs --- a/src/Futhark/Pipeline.hs +++ b/src/Futhark/Pipeline.hs @@ -22,6 +22,7 @@ module Futhark.Error, onePass, passes, + condPipeline, runPipeline, ) where @@ -36,6 +37,7 @@ import qualified Data.Text.IO as T import Data.Time.Clock import qualified Futhark.Analysis.Alias as Alias +import Futhark.Compiler.Config (Verbosity (..)) import Futhark.Error import Futhark.IR (PrettyRep, Prog) import Futhark.IR.TypeCheck @@ -47,16 +49,6 @@ import Text.Printf import Prelude hiding (id, (.)) --- | 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) - newtype FutharkEnv = FutharkEnv {futharkVerbose :: Verbosity} data FutharkState = FutharkState @@ -153,6 +145,15 @@ Left err -> validationError pass prog'' $ show err Right () -> pure () pure prog' + +-- | Conditionally run pipeline if predicate is true. +condPipeline :: + (Prog rep -> Bool) -> Pipeline rep rep -> Pipeline rep rep +condPipeline cond (Pipeline f) = + Pipeline $ \cfg prog -> + if cond prog + then f cfg prog + else pure prog -- | Create a pipeline from a list of passes. passes :: diff --git a/src/Futhark/Script.hs b/src/Futhark/Script.hs --- a/src/Futhark/Script.hs +++ b/src/Futhark/Script.hs @@ -152,8 +152,8 @@ [ try $ inParens sep (mkTuple <$> (parseExp sep `sepBy` pComma)), inParens sep $ parseExp sep, inBraces sep (Record <$> (pField `sepBy` pComma)), - Const <$> V.parseValue sep, StringLit . T.pack <$> lexeme sep ("\"" *> manyTill charLiteral "\""), + Const <$> V.parseValue sep, Call <$> parseFunc <*> pure [] ] @@ -378,7 +378,7 @@ evalExp' _ (StringLit s) = case V.putValue s of Just s' -> - pure $ V.ValueAtom $ SValue (V.valueTypeText (V.valueType s')) $ VVal s' + pure $ V.ValueAtom $ SValue (V.valueTypeTextNoDims (V.valueType s')) $ VVal s' Nothing -> error $ "Unable to write value " ++ pretty s evalExp' _ (Const val) = pure $ V.ValueAtom $ SValue (V.valueTypeTextNoDims (V.valueType val)) $ VVal val diff --git a/src/Futhark/Test.hs b/src/Futhark/Test.hs --- a/src/Futhark/Test.hs +++ b/src/Futhark/Test.hs @@ -163,11 +163,15 @@ note (_, t) | "(" `T.isPrefixOf` t = Just $ - "\nNote: expected type " <> prettyText t <> " is an opaque tuple that cannot be constructed\n" + "\nNote: expected type " + <> prettyText t + <> " is an opaque tuple that cannot be constructed\n" <> "in FutharkScript. Consider using type annotations to give it a proper name." | "{" `T.isPrefixOf` t = Just $ - "\nNote: expected type " <> prettyText t <> " is an opaque record that cannot be constructed\n" + "\nNote: expected type " + <> prettyText t + <> " is an opaque record that cannot be constructed\n" <> "in FutharkScript. Consider using type annotations to give it a proper name." | otherwise = Nothing diff --git a/src/Futhark/Transform/FirstOrderTransform.hs b/src/Futhark/Transform/FirstOrderTransform.hs --- a/src/Futhark/Transform/FirstOrderTransform.hs +++ b/src/Futhark/Transform/FirstOrderTransform.hs @@ -123,6 +123,10 @@ Pat (LetDec (Rep m)) -> SOAC (Rep m) -> m () +transformSOAC _ JVP {} = + error "transformSOAC: unhandled JVP" +transformSOAC _ VJP {} = + error "transformSOAC: unhandled VJP" transformSOAC pat (Screma w arrs form@(ScremaForm scans reds map_lam)) = do -- See Note [Translation of Screma]. -- diff --git a/src/Futhark/Util.hs b/src/Futhark/Util.hs --- a/src/Futhark/Util.hs +++ b/src/Futhark/Util.hs @@ -15,6 +15,8 @@ maxinum, chunk, chunks, + pairs, + unpairs, dropAt, takeLast, dropLast, @@ -55,13 +57,13 @@ trim, pmapIO, interactWithFileSafely, - readFileSafely, convFloat, UserString, EncodedString, zEncodeString, atMostChars, invertMap, + traverseFold, fixPoint, ) where @@ -76,6 +78,7 @@ import qualified Data.ByteString.Base16 as Base16 import Data.Char import Data.Either +import Data.Foldable (fold) import Data.Function ((&)) import Data.List (foldl', genericDrop, genericSplitAt, sortBy) import qualified Data.List.NonEmpty as NE @@ -85,7 +88,6 @@ import qualified Data.Text as T import qualified Data.Text.Encoding as T import qualified Data.Text.Encoding.Error as T -import qualified Data.Text.IO as T import Data.Time.Clock (UTCTime, getCurrentTime) import Data.Tuple (swap) import Numeric @@ -142,6 +144,18 @@ let (bef, aft) = splitAt n xs in bef : chunks ns aft +-- | @pairs l@ chunks the list into pairs of consecutive elements, +-- ignoring any excess element. Example: @pairs [a,b,c,d] == +-- [(a,b),(c,d)]@. +pairs :: [a] -> [(a, a)] +pairs (a : b : l) = (a, b) : pairs l +pairs _ = [] + +-- | The opposite of 'pairs': @unpairs [(a,b),(c,d)] = [a,b,c,d]@. +unpairs :: [(a, a)] -> [a] +unpairs [] = [] +unpairs ((a, b) : l) = a : b : unpairs l + -- | Like 'maximum', but returns zero for an empty list. maxinum :: (Num a, Ord a, Foldable f) => f a -> a maxinum = foldl' max 0 @@ -422,12 +436,6 @@ | otherwise = pure $ Just $ Left $ show e --- | Read a file, returning 'Nothing' if the file does not exist, and --- 'Left' if some other error occurs. -readFileSafely :: FilePath -> IO (Maybe (Either String T.Text)) -readFileSafely filepath = - interactWithFileSafely $ T.readFile filepath - -- | Convert between different floating-point types, preserving -- infinities and NaNs. convFloat :: (RealFloat from, RealFloat to) => from -> to @@ -528,6 +536,10 @@ M.toList m & fmap (swap . first S.singleton) & foldr (uncurry $ M.insertWith (<>)) mempty + +-- | Applicatively fold a traversable. +traverseFold :: (Monoid m, Traversable t, Applicative f) => (a -> f m) -> t a -> f m +traverseFold f = fmap fold . traverse f -- | Perform fixpoint iteration. fixPoint :: Eq a => (a -> a) -> a -> a diff --git a/src/Futhark/Util/Console.hs b/src/Futhark/Util/Console.hs --- a/src/Futhark/Util/Console.hs +++ b/src/Futhark/Util/Console.hs @@ -2,6 +2,7 @@ module Futhark.Util.Console ( color, inRed, + inYellow, inBold, ) where @@ -15,6 +16,10 @@ -- | Make the string red. inRed :: String -> String inRed s = setSGRCode [SetColor Foreground Vivid Red] ++ s ++ setSGRCode [Reset] + +-- | Make the string yellow. +inYellow :: String -> String +inYellow s = setSGRCode [SetColor Foreground Vivid Yellow] ++ s ++ setSGRCode [Reset] -- | Make the string bold. inBold :: String -> String diff --git a/src/Futhark/Util/IntegralExp.hs b/src/Futhark/Util/IntegralExp.hs --- a/src/Futhark/Util/IntegralExp.hs +++ b/src/Futhark/Util/IntegralExp.hs @@ -31,6 +31,7 @@ div :: e -> e -> e mod :: e -> e -> e sgn :: e -> Maybe Int + pow :: e -> e -> e -- | Like 'Futhark.Util.IntegralExp.div', but rounds towards -- positive infinity. @@ -71,3 +72,4 @@ div = liftOp2 Prelude.div mod = liftOp2 Prelude.mod sgn = Just . fromIntegral . signum . toInteger . wrappedValue + pow = liftOp2 (Prelude.^) diff --git a/src/Futhark/Util/ProgressBar.hs b/src/Futhark/Util/ProgressBar.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Util/ProgressBar.hs @@ -0,0 +1,57 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | Facilities for generating and otherwise handling text-based progress bars. +module Futhark.Util.ProgressBar + ( progressBar, + ProgressBar (..), + progressSpinner, + ) +where + +import qualified Data.Text as T + +-- | Information about a progress bar to render. The "progress space" +-- spans from 0 and up to the `progressBarBound`, but can be +-- visualised in any number of steps. +data ProgressBar = ProgressBar + { -- | Number of steps in the visualisation. + progressBarSteps :: Int, + -- | The logical upper bound. + progressBarBound :: Double, + -- | The current position in the progress bar, relative to the + -- upper bound. + progressBarElapsed :: Double + } + +-- | Render the progress bar. +progressBar :: ProgressBar -> T.Text +progressBar (ProgressBar steps bound elapsed) = + "|" <> T.pack (map cell [1 .. steps]) <> "| " + where + step_size :: Double + step_size = bound / fromIntegral steps + chars = " ▏▎▍▍▌▋▊▉█" + num_chars = T.length chars + char i + | i >= 0 && i < num_chars = T.index chars i + | otherwise = ' ' + + cell :: Int -> Char + cell i + | i' * step_size <= elapsed = char 9 + | otherwise = + char (floor (((elapsed - (i' - 1) * step_size) * fromIntegral num_chars) / step_size)) + where + i' = fromIntegral i + +-- | Render a spinner - a kind of progress bar where there is no upper +-- bound because we don't know how long it'll take. You certainly +-- know these from THE INTERNET. The non-negative integer is how many +-- "steps" have been taken. The spinner looks best if this is +-- incremented by one for every call. +progressSpinner :: Int -> T.Text +progressSpinner spin_idx = + T.singleton $ T.index spin_load (spin_idx `rem` n) + where + spin_load = "⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏" + n = T.length spin_load diff --git a/src/Language/Futhark/Interpreter.hs b/src/Language/Futhark/Interpreter.hs --- a/src/Language/Futhark/Interpreter.hs +++ b/src/Language/Futhark/Interpreter.hs @@ -1933,6 +1933,10 @@ <> pretty (asInt64 m) <> "]" else pure $ toArray shape $ map (toArray rowshape) $ chunk (asInt m) xs' + def "vjp2" = Just $ + fun3t $ \_ _ _ -> bad noLoc mempty "Interpreter does not support autodiff." + def "jvp2" = Just $ + fun3t $ \_ _ _ -> bad noLoc mempty "Interpreter does not support autodiff." def "acc" = Nothing def s | nameFromString s `M.member` namesToPrimTypes = Nothing def s = error $ "Missing intrinsic: " ++ s diff --git a/src/Language/Futhark/Prop.hs b/src/Language/Futhark/Prop.hs --- a/src/Language/Futhark/Prop.hs +++ b/src/Language/Futhark/Prop.hs @@ -73,6 +73,7 @@ DimPos (..), mustBeExplicit, mustBeExplicitInType, + determineSizeWitnesses, tupleRecord, isTupleRecord, areTupleFields, @@ -234,13 +235,22 @@ onDim _ _ _ = pure () +-- | Determine which of the sizes in a type are used as sizes outside +-- of functions in the type, and which are not. The former are said +-- to be "witnessed" by this type, while the latter are not. In +-- practice, the latter means that the actual sizes must come from +-- somewhere else. +determineSizeWitnesses :: StructType -> (S.Set VName, S.Set VName) +determineSizeWitnesses t = + bimap (S.fromList . M.keys) (S.fromList . M.keys) $ + M.partition not $ mustBeExplicitAux t + -- | Figure out which of the sizes in a parameter type must be passed -- explicitly, because their first use is as something else than just -- an array dimension. 'mustBeExplicit' is like this function, but -- first decomposes into parameter types. mustBeExplicitInType :: StructType -> S.Set VName -mustBeExplicitInType t = - S.fromList $ M.keys $ M.filter id $ mustBeExplicitAux t +mustBeExplicitInType = snd . determineSizeWitnesses -- | Figure out which of the sizes in a binding type must be passed -- explicitly, because their first use is as something else than just @@ -1068,6 +1078,24 @@ arr_b $ shape [n] ] $ RetType [] $ uarr_a $ shape [k] + ), + ( "jvp2", + IntrinsicPolyFun + [tp_a, tp_b] + [ Scalar t_a `arr` Scalar t_b, + Scalar t_a, + Scalar t_a + ] + $ RetType [] $ Scalar $ tupleRecord [Scalar t_b, Scalar t_b] + ), + ( "vjp2", + IntrinsicPolyFun + [tp_a, tp_b] + [ Scalar t_a `arr` Scalar t_b, + Scalar t_a, + Scalar t_b + ] + $ RetType [] $ Scalar $ tupleRecord [Scalar t_b, Scalar t_a] ) ] ++ @@ -1315,34 +1343,62 @@ -- | The set of module types used in any exported (non-local) -- declaration. -progModuleTypes :: Ord vn => ProgBase f vn -> S.Set vn -progModuleTypes = mconcat . map onDec . progDecs +progModuleTypes :: ProgBase Info VName -> S.Set VName +progModuleTypes prog = foldMap reach mtypes_used where - onDec (OpenDec x _) = onModExp x - onDec (ModDec md) = - maybe mempty (onSigExp . fst) (modSignature md) <> onModExp (modExp md) - onDec SigDec {} = mempty - onDec TypeDec {} = mempty - onDec ValDec {} = mempty - onDec LocalDec {} = mempty - onDec ImportDec {} = mempty + -- Fixed point iteration. + reach v = S.singleton v <> maybe mempty (foldMap reach) (M.lookup v reachable_from_mtype) - onModExp ModVar {} = mempty - onModExp (ModParens p _) = onModExp p - onModExp ModImport {} = mempty - onModExp (ModDecs ds _) = mconcat $ map onDec ds - onModExp (ModApply me1 me2 _ _ _) = onModExp me1 <> onModExp me2 - onModExp (ModAscript me se _ _) = onModExp me <> onSigExp se - onModExp (ModLambda p r me _) = - onModParam p <> maybe mempty (onSigExp . fst) r <> onModExp me + reachable_from_mtype = foldMap onDec $ progDecs prog + where + onDec OpenDec {} = mempty + onDec ModDec {} = mempty + onDec (SigDec sb) = + M.singleton (sigName sb) (onSigExp (sigExp sb)) + onDec TypeDec {} = mempty + onDec ValDec {} = mempty + onDec (LocalDec d _) = onDec d + onDec ImportDec {} = mempty - onModParam = onSigExp . modParamType + onSigExp (SigVar v _ _) = S.singleton $ qualLeaf v + onSigExp (SigParens e _) = onSigExp e + onSigExp (SigSpecs ss _) = foldMap onSpec ss + onSigExp (SigWith e _ _) = onSigExp e + onSigExp (SigArrow _ e1 e2 _) = onSigExp e1 <> onSigExp e2 - onSigExp (SigVar v _ _) = S.singleton $ qualLeaf v - onSigExp (SigParens e _) = onSigExp e - onSigExp SigSpecs {} = mempty - onSigExp (SigWith e _ _) = onSigExp e - onSigExp (SigArrow _ e1 e2 _) = onSigExp e1 <> onSigExp e2 + onSpec ValSpec {} = mempty + onSpec TypeSpec {} = mempty + onSpec TypeAbbrSpec {} = mempty + onSpec (ModSpec vn e _ _) = S.singleton vn <> onSigExp e + onSpec (IncludeSpec e _) = onSigExp e + + mtypes_used = foldMap onDec $ progDecs prog + where + onDec (OpenDec x _) = onModExp x + onDec (ModDec md) = + maybe mempty (onSigExp . fst) (modSignature md) <> onModExp (modExp md) + onDec SigDec {} = mempty + onDec TypeDec {} = mempty + onDec ValDec {} = mempty + onDec LocalDec {} = mempty + onDec ImportDec {} = mempty + + onModExp ModVar {} = mempty + onModExp (ModParens p _) = onModExp p + onModExp ModImport {} = mempty + onModExp (ModDecs ds _) = mconcat $ map onDec ds + onModExp (ModApply me1 me2 _ _ _) = onModExp me1 <> onModExp me2 + onModExp (ModAscript me se _ _) = onModExp me <> onSigExp se + onModExp (ModLambda p r me _) = + onModParam p <> maybe mempty (onSigExp . fst) r <> onModExp me + + onModParam = onSigExp . modParamType + + onSigExp (SigVar v _ _) = S.singleton $ qualLeaf v + onSigExp (SigParens e _) = onSigExp e + onSigExp SigSpecs {} = mempty + onSigExp (SigWith e _ _) = onSigExp e + onSigExp (SigArrow _ e1 e2 _) = onSigExp e1 <> onSigExp e2 -- | Extract a leading @((name, namespace, file), remainder)@ from a -- documentation comment string. These are formatted as diff --git a/src/Language/Futhark/Query.hs b/src/Language/Futhark/Query.hs --- a/src/Language/Futhark/Query.hs +++ b/src/Language/Futhark/Query.hs @@ -371,10 +371,7 @@ containingModule imports (Pos file _ _ _) = snd <$> find ((== file') . fst) imports where - file' = - includeToString $ - mkInitialImport $ - fst $ Posix.splitExtension file + file' = includeToString $ mkInitialImport $ fst $ Posix.splitExtension file -- | Information about what is at the given source location. data AtPos = AtName (QualName VName) (Maybe BoundTo) Loc diff --git a/src/Language/Futhark/Syntax.hs b/src/Language/Futhark/Syntax.hs --- a/src/Language/Futhark/Syntax.hs +++ b/src/Language/Futhark/Syntax.hs @@ -92,7 +92,6 @@ DecBase (..), -- * Miscellaneous - Showable, NoInfo (..), Info (..), Alias (..), @@ -125,38 +124,11 @@ import Language.Futhark.Core import Prelude --- | Convenience class for deriving 'Show' instances for the AST. -class - ( Show vn, - Show (f VName), - Show (f (Diet, Maybe VName)), - Show (f String), - Show (f [VName]), - Show (f ([VName], [VName])), - Show (f PatType), - Show (f (PatType, [VName])), - Show (f (StructType, [VName])), - Show (f (StructRetType, [VName])), - Show (f EntryPoint), - Show (f StructType), - Show (f StructRetType), - Show (f PatRetType), - Show (f (StructType, Maybe VName)), - Show (f (PName, StructType)), - Show (f (PName, StructType, Maybe VName)), - Show (f (Aliasing, StructRetType)), - Show (f (M.Map VName VName)), - Show (f AppRes) - ) => - Showable f vn - -- | No information functor. Usually used for placeholder type- or -- aliasing information. data NoInfo a = NoInfo deriving (Eq, Ord, Show) -instance Show vn => Showable NoInfo vn - instance Functor NoInfo where fmap _ NoInfo = NoInfo @@ -170,8 +142,6 @@ newtype Info a = Info {unInfo :: a} deriving (Eq, Ord, Show) -instance Show vn => Showable Info vn - instance Functor Info where fmap f (Info x) = Info $ f x @@ -267,12 +237,10 @@ NamedDim (QualName vn) | -- | The size is a constant. ConstDim Int - | -- | No known size - but still possibly given a unique name, so we - -- can recognise e.g. @type square [n] = [n][n]i32@ and make - -- @square []@ do the right thing. If @Nothing@, then this is a - -- name distinct from any other. The type checker should _never_ - -- produce these - they are a (hopefully temporary) thing - -- introduced by defunctorisation and monomorphisation. + | -- | No known size. If @Nothing@, then this is a name distinct + -- from any other. The type checker should _never_ produce these + -- - they are a (hopefully temporary) thing introduced by + -- defunctorisation and monomorphisation. AnyDim (Maybe vn) deriving (Show) @@ -553,8 +521,10 @@ expandedType :: f StructType } -deriving instance Showable f vn => Show (TypeDeclBase f vn) +deriving instance Show (TypeDeclBase Info VName) +deriving instance Show vn => Show (TypeDeclBase NoInfo vn) + deriving instance Eq (TypeDeclBase NoInfo VName) deriving instance Ord (TypeDeclBase NoInfo VName) @@ -596,8 +566,10 @@ identSrcLoc :: SrcLoc } -deriving instance Showable f vn => Show (IdentBase f vn) +deriving instance Show (IdentBase Info VName) +deriving instance Show vn => Show (IdentBase NoInfo vn) + instance Eq vn => Eq (IdentBase ty vn) where x == y = identName x == identName y @@ -678,8 +650,10 @@ (Maybe (ExpBase f vn)) (Maybe (ExpBase f vn)) -deriving instance Showable f vn => Show (DimIndexBase f vn) +deriving instance Show (DimIndexBase Info VName) +deriving instance Show vn => Show (DimIndexBase NoInfo vn) + deriving instance Eq (DimIndexBase NoInfo VName) deriving instance Ord (DimIndexBase NoInfo VName) @@ -789,8 +763,10 @@ | -- | A match expression. Match (ExpBase f vn) (NE.NonEmpty (CaseBase f vn)) SrcLoc -deriving instance Showable f vn => Show (AppExpBase f vn) +deriving instance Show (AppExpBase Info VName) +deriving instance Show vn => Show (AppExpBase NoInfo vn) + deriving instance Eq (AppExpBase NoInfo VName) deriving instance Ord (AppExpBase NoInfo VName) @@ -893,8 +869,10 @@ Ascript (ExpBase f vn) (TypeDeclBase f vn) SrcLoc | AppExp (AppExpBase f vn) (f AppRes) -deriving instance Showable f vn => Show (ExpBase f vn) +deriving instance Show (ExpBase Info VName) +deriving instance Show vn => Show (ExpBase NoInfo vn) + deriving instance Eq (ExpBase NoInfo VName) deriving instance Ord (ExpBase NoInfo VName) @@ -933,8 +911,10 @@ = RecordFieldExplicit Name (ExpBase f vn) SrcLoc | RecordFieldImplicit vn (f PatType) SrcLoc -deriving instance Showable f vn => Show (FieldBase f vn) +deriving instance Show (FieldBase Info VName) +deriving instance Show vn => Show (FieldBase NoInfo vn) + deriving instance Eq (FieldBase NoInfo VName) deriving instance Ord (FieldBase NoInfo VName) @@ -946,8 +926,10 @@ -- | A case in a match expression. data CaseBase f vn = CasePat (PatBase f vn) (ExpBase f vn) SrcLoc -deriving instance Showable f vn => Show (CaseBase f vn) +deriving instance Show (CaseBase Info VName) +deriving instance Show vn => Show (CaseBase NoInfo vn) + deriving instance Eq (CaseBase NoInfo VName) deriving instance Ord (CaseBase NoInfo VName) @@ -961,8 +943,10 @@ | ForIn (PatBase f vn) (ExpBase f vn) | While (ExpBase f vn) -deriving instance Showable f vn => Show (LoopFormBase f vn) +deriving instance Show (LoopFormBase Info VName) +deriving instance Show vn => Show (LoopFormBase NoInfo vn) + deriving instance Eq (LoopFormBase NoInfo VName) deriving instance Ord (LoopFormBase NoInfo VName) @@ -987,8 +971,10 @@ | PatConstr Name (f PatType) [PatBase f vn] SrcLoc | PatAttr (AttrInfo vn) (PatBase f vn) SrcLoc -deriving instance Showable f vn => Show (PatBase f vn) +deriving instance Show (PatBase Info VName) +deriving instance Show vn => Show (PatBase NoInfo vn) + deriving instance Eq (PatBase NoInfo VName) deriving instance Ord (PatBase NoInfo VName) @@ -1057,8 +1043,10 @@ valBindLocation :: SrcLoc } -deriving instance Showable f vn => Show (ValBindBase f vn) +deriving instance Show (ValBindBase Info VName) +deriving instance Show (ValBindBase NoInfo Name) + instance Located (ValBindBase f vn) where locOf = locOf . valBindLocation @@ -1073,8 +1061,10 @@ typeBindLocation :: SrcLoc } -deriving instance Showable f vn => Show (TypeBindBase f vn) +deriving instance Show (TypeBindBase Info VName) +deriving instance Show (TypeBindBase NoInfo Name) + instance Located (TypeBindBase f vn) where locOf = locOf . typeBindLocation @@ -1133,8 +1123,10 @@ | ModSpec vn (SigExpBase f vn) (Maybe DocComment) SrcLoc | IncludeSpec (SigExpBase f vn) SrcLoc -deriving instance Showable f vn => Show (SpecBase f vn) +deriving instance Show (SpecBase Info VName) +deriving instance Show (SpecBase NoInfo Name) + instance Located (SpecBase f vn) where locOf (ValSpec _ _ _ _ loc) = locOf loc locOf (TypeAbbrSpec tbind) = locOf tbind @@ -1150,13 +1142,17 @@ | SigWith (SigExpBase f vn) (TypeRefBase f vn) SrcLoc | SigArrow (Maybe vn) (SigExpBase f vn) (SigExpBase f vn) SrcLoc -deriving instance Showable f vn => Show (SigExpBase f vn) +deriving instance Show (SigExpBase Info VName) +deriving instance Show (SigExpBase NoInfo Name) + -- | A type refinement. data TypeRefBase f vn = TypeRef (QualName vn) [TypeParamBase vn] (TypeDeclBase f vn) SrcLoc -deriving instance Showable f vn => Show (TypeRefBase f vn) +deriving instance Show (TypeRefBase Info VName) +deriving instance Show (TypeRefBase NoInfo Name) + instance Located (TypeRefBase f vn) where locOf (TypeRef _ _ _ loc) = locOf loc @@ -1175,8 +1171,10 @@ sigLoc :: SrcLoc } -deriving instance Showable f vn => Show (SigBindBase f vn) +deriving instance Show (SigBindBase Info VName) +deriving instance Show (SigBindBase NoInfo Name) + instance Located (SigBindBase f vn) where locOf = locOf . sigLoc @@ -1203,8 +1201,10 @@ (ModExpBase f vn) SrcLoc -deriving instance Showable f vn => Show (ModExpBase f vn) +deriving instance Show (ModExpBase Info VName) +deriving instance Show (ModExpBase NoInfo Name) + instance Located (ModExpBase f vn) where locOf (ModVar _ loc) = locOf loc locOf (ModParens _ loc) = locOf loc @@ -1224,8 +1224,10 @@ modLocation :: SrcLoc } -deriving instance Showable f vn => Show (ModBindBase f vn) +deriving instance Show (ModBindBase Info VName) +deriving instance Show (ModBindBase NoInfo Name) + instance Located (ModBindBase f vn) where locOf = locOf . modLocation @@ -1237,8 +1239,10 @@ modParamLocation :: SrcLoc } -deriving instance Showable f vn => Show (ModParamBase f vn) +deriving instance Show (ModParamBase Info VName) +deriving instance Show (ModParamBase NoInfo Name) + instance Located (ModParamBase f vn) where locOf = locOf . modParamLocation @@ -1252,8 +1256,10 @@ | LocalDec (DecBase f vn) SrcLoc | ImportDec FilePath (f FilePath) SrcLoc -deriving instance Showable f vn => Show (DecBase f vn) +deriving instance Show (DecBase Info VName) +deriving instance Show (DecBase NoInfo Name) + instance Located (DecBase f vn) where locOf (ValDec d) = locOf d locOf (TypeDec d) = locOf d @@ -1270,7 +1276,9 @@ progDecs :: [DecBase f vn] } -deriving instance Showable f vn => Show (ProgBase f vn) +deriving instance Show (ProgBase Info VName) + +deriving instance Show (ProgBase NoInfo Name) --- Some prettyprinting definitions are here because we need them in --- the Attributes module. diff --git a/src/Language/Futhark/TypeChecker/Modules.hs b/src/Language/Futhark/TypeChecker/Modules.hs --- a/src/Language/Futhark/TypeChecker/Modules.hs +++ b/src/Language/Futhark/TypeChecker/Modules.hs @@ -266,13 +266,13 @@ mismatchedLiftedness name_l (map qualLeaf $ M.keys mod_abs) - (qualLeaf name) + name (mod_l, ps, t) | name_l < SizeLifted, not $ null $ retDims t -> anonymousSizes (map qualLeaf $ M.keys mod_abs) - (qualLeaf name) + name (mod_l, ps, t) | Just (abs_name, _) <- M.lookup (fmap baseName name) abs_mapping -> pure (qualLeaf name, (abs_name, TypeAbbr name_l ps t)) @@ -355,25 +355,26 @@ mismatchedType :: Loc -> [VName] -> + [VName] -> VName -> (Liftedness, [TypeParam], StructRetType) -> (Liftedness, [TypeParam], StructRetType) -> Either TypeError b -mismatchedType loc abs name spec_t env_t = +mismatchedType loc abs quals name spec_t env_t = Left . TypeError loc mempty $ "Module defines" - indent 2 (ppTypeAbbr abs name env_t) + indent 2 (ppTypeAbbr abs (QualName quals name) env_t) "but module type requires" - indent 2 (ppTypeAbbr abs name spec_t) + indent 2 (ppTypeAbbr abs (QualName quals name) spec_t) -ppTypeAbbr :: [VName] -> VName -> (Liftedness, [TypeParam], StructRetType) -> Doc +ppTypeAbbr :: [VName] -> QualName VName -> (Liftedness, [TypeParam], StructRetType) -> Doc ppTypeAbbr abs name (l, ps, RetType [] (Scalar (TypeVar () _ tn args))) | typeLeaf tn `elem` abs, map typeParamToArg ps == args = - "type" <> ppr l <+> pprName name + "type" <> ppr l <+> ppr name <+> spread (map ppr ps) ppTypeAbbr _ name (l, ps, t) = - "type" <> ppr l <+> pprName name + "type" <> ppr l <+> ppr name <+> spread (map ppr ps) <+> equals <+/> nest 2 (align (ppr t)) @@ -390,29 +391,31 @@ matchMTys orig_mty orig_mty_sig = matchMTys' (M.map (SizeSubst . NamedDim) $ resolveMTyNames orig_mty orig_mty_sig) + [] orig_mty orig_mty_sig where matchMTys' :: M.Map VName (Subst StructRetType) -> + [VName] -> MTy -> MTy -> Loc -> Either TypeError (M.Map VName VName) - matchMTys' _ (MTy _ ModFun {}) (MTy _ ModEnv {}) loc = + matchMTys' _ _ (MTy _ ModFun {}) (MTy _ ModEnv {}) loc = Left $ TypeError loc mempty "Cannot match parametric module with non-parametric module type." - matchMTys' _ (MTy _ ModEnv {}) (MTy _ ModFun {}) loc = + matchMTys' _ _ (MTy _ ModEnv {}) (MTy _ ModFun {}) loc = Left $ TypeError loc mempty "Cannot match non-parametric module with paramatric module type." - matchMTys' old_abs_subst_to_type (MTy mod_abs mod) (MTy sig_abs sig) loc = do + matchMTys' old_abs_subst_to_type quals (MTy mod_abs mod) (MTy sig_abs sig) loc = do -- Check that abstract types in 'sig' have an implementation in -- 'mod'. This also gives us a substitution that we use to check -- the types of values. @@ -421,31 +424,33 @@ let abs_subst_to_type = old_abs_subst_to_type <> M.map (substFromAbbr . snd) abs_substs abs_name_substs = M.map (qualLeaf . fst) abs_substs - substs <- matchMods abs_subst_to_type mod sig loc + substs <- matchMods abs_subst_to_type quals mod sig loc pure (substs <> abs_name_substs) matchMods :: M.Map VName (Subst StructRetType) -> + [VName] -> Mod -> Mod -> Loc -> Either TypeError (M.Map VName VName) - matchMods _ ModEnv {} ModFun {} loc = + matchMods _ _ ModEnv {} ModFun {} loc = Left $ TypeError loc mempty "Cannot match non-parametric module with parametric module type." - matchMods _ ModFun {} ModEnv {} loc = + matchMods _ _ ModFun {} ModEnv {} loc = Left $ TypeError loc mempty "Cannot match parametric module with non-parametric module type." - matchMods abs_subst_to_type (ModEnv mod) (ModEnv sig) loc = - matchEnvs abs_subst_to_type mod sig loc + matchMods abs_subst_to_type quals (ModEnv mod) (ModEnv sig) loc = + matchEnvs abs_subst_to_type quals mod sig loc matchMods old_abs_subst_to_type + quals (ModFun (FunSig mod_abs mod_pmod mod_mod)) (ModFun (FunSig sig_abs sig_pmod sig_mod)) loc = do @@ -453,17 +458,18 @@ let abs_subst_to_type = old_abs_subst_to_type <> M.map (substFromAbbr . snd) abs_substs abs_name_substs = M.map (qualLeaf . fst) abs_substs - pmod_substs <- matchMods abs_subst_to_type mod_pmod sig_pmod loc - mod_substs <- matchMTys' abs_subst_to_type mod_mod sig_mod loc + pmod_substs <- matchMods abs_subst_to_type quals mod_pmod sig_pmod loc + mod_substs <- matchMTys' abs_subst_to_type quals mod_mod sig_mod loc pure (pmod_substs <> mod_substs <> abs_name_substs) matchEnvs :: M.Map VName (Subst StructRetType) -> + [VName] -> Env -> Env -> Loc -> Either TypeError (M.Map VName VName) - matchEnvs abs_subst_to_type env sig loc = do + matchEnvs abs_subst_to_type quals env sig loc = do -- XXX: we only want to create substitutions for visible names. -- This must be wrong in some cases. Probably we need to -- rethink how we do shadowing for module types. @@ -476,7 +482,7 @@ \(name, TypeAbbr spec_l spec_ps spec_t) -> case findBinding envTypeTable Type (baseName name) env of Just (name', TypeAbbr l ps t) -> - matchTypeAbbr loc abs_subst_to_type name spec_l spec_ps spec_t name' l ps t + matchTypeAbbr loc abs_subst_to_type quals name spec_l spec_ps spec_t name' l ps t Nothing -> missingType loc $ baseName name -- Check that all values are defined correctly, substituting the @@ -485,7 +491,7 @@ forM (M.toList $ envVtable sig) $ \(name, spec_bv) -> do let spec_bv' = substituteTypesInBoundV (`M.lookup` abs_subst_to_type) spec_bv case findBinding envVtable Term (baseName name) env of - Just (name', bv) -> matchVal loc name spec_bv' name' bv + Just (name', bv) -> matchVal loc quals name spec_bv' name' bv _ -> missingVal loc (baseName name) -- Check for correct modules. @@ -493,7 +499,8 @@ forM (M.toList $ envModTable sig) $ \(name, modspec) -> case findBinding envModTable Term (baseName name) env of Just (name', mod) -> - M.insert name name' <$> matchMods abs_subst_to_type mod modspec loc + M.insert name name' + <$> matchMods abs_subst_to_type (quals ++ [name]) mod modspec loc Nothing -> missingMod loc $ baseName name @@ -502,6 +509,7 @@ matchTypeAbbr :: Loc -> M.Map VName (Subst StructRetType) -> + [VName] -> VName -> Liftedness -> [TypeParam] -> @@ -511,11 +519,9 @@ [TypeParam] -> StructRetType -> Either TypeError (VName, VName) - matchTypeAbbr loc abs_subst_to_type spec_name spec_l spec_ps spec_t name l ps t = do - -- We have to create substitutions for the type parameters, too. + matchTypeAbbr loc abs_subst_to_type quals spec_name spec_l spec_ps spec_t name l ps t = do + -- Number of type parameters must match. unless (length spec_ps == length ps) $ nomatch spec_t - param_substs <- - mconcat <$> zipWithM (matchTypeParam (nomatch spec_t)) spec_ps ps -- Abstract types have a particular restriction to ensure that -- if we have a value of an abstract type 't [n]', then there is @@ -526,67 +532,60 @@ d : _ -> Left . TypeError loc mempty $ "Type" - indent 2 (ppTypeAbbr [] name (l, ps, t)) + indent 2 (ppTypeAbbr [] (QualName quals name) (l, ps, t)) textwrap "cannot be made abstract because size parameter" <+/> pquote (pprName d) <+/> textwrap "is not used as an array size in the definition." - let spec_t' = applySubst (`M.lookup` (param_substs <> abs_subst_to_type)) spec_t - if spec_t' == t - then pure (spec_name, name) - else nomatch spec_t' + let spec_t' = applySubst (`M.lookup` abs_subst_to_type) spec_t + nonrigid = ps <> map (`TypeParamDim` mempty) (retDims t) + case doUnification loc spec_ps nonrigid (retType spec_t') (retType t) of + Right t' + | noSizes t' `subtypeOf` noSizes (retType spec_t') -> pure (spec_name, name) + _ -> nomatch spec_t' where nomatch spec_t' = mismatchedType loc (M.keys abs_subst_to_type) + quals spec_name (spec_l, spec_ps, spec_t') (l, ps, t) - matchTypeParam _ (TypeParamDim x _) (TypeParamDim y _) = - pure $ M.singleton x $ SizeSubst $ NamedDim $ qualName y - matchTypeParam _ (TypeParamType spec_l x _) (TypeParamType l y _) - | spec_l <= l = - pure . M.singleton x . Subst [] . RetType [] $ - Scalar $ TypeVar () Nonunique (typeName y) [] - matchTypeParam nomatch _ _ = - nomatch - matchVal :: Loc -> + [VName] -> VName -> BoundV -> VName -> BoundV -> Either TypeError (VName, VName) - matchVal loc spec_name spec_v name v = + matchVal loc quals spec_name spec_v name v = case matchValBinding loc spec_v v of Nothing -> pure (spec_name, name) Just problem -> Left $ TypeError loc mempty $ "Module type specifies" - indent 2 (ppValBind spec_name spec_v) + indent 2 (ppValBind (QualName quals spec_name) spec_v) "but module provides" - indent 2 (ppValBind spec_name v) + indent 2 (ppValBind (QualName quals spec_name) v) fromMaybe mempty problem matchValBinding :: Loc -> BoundV -> BoundV -> Maybe (Maybe Doc) - matchValBinding loc (BoundV _ orig_spec_t) (BoundV tps orig_t) = - case doUnification loc tps (toStruct orig_spec_t) (toStruct orig_t) of + matchValBinding loc (BoundV spec_tps orig_spec_t) (BoundV tps orig_t) = do + case doUnification loc spec_tps tps (toStruct orig_spec_t) (toStruct orig_t) of Left (TypeError _ notes msg) -> Just $ Just $ msg <> ppr notes -- Even if they unify, we still have to verify the uniqueness -- properties. Right t - | noSizes t - `subtypeOf` noSizes orig_spec_t -> - Nothing + | noSizes t `subtypeOf` noSizes orig_spec_t -> Nothing | otherwise -> Just Nothing ppValBind v (BoundV tps t) = - "val" <+> pprName v <+> spread (map ppr tps) <+> colon + "val" <+> ppr v <+> spread (map ppr tps) <+> colon indent 2 (align (ppr t)) -- | Apply a parametric module to an argument. diff --git a/src/Language/Futhark/TypeChecker/Monad.hs b/src/Language/Futhark/TypeChecker/Monad.hs --- a/src/Language/Futhark/TypeChecker/Monad.hs +++ b/src/Language/Futhark/TypeChecker/Monad.hs @@ -287,7 +287,8 @@ Scalar (Prim (Signed Int64)) -> pure v' _ -> typeError loc mempty $ - "Dimension declaration" <+> ppr v <+> "should be of type i64." + "Sizes must have type i64, but" <+> pquote (ppr v) <+> "has type:" + ppr t typeError :: Located loc => loc -> Notes -> Doc -> m a diff --git a/src/Language/Futhark/TypeChecker/Terms.hs b/src/Language/Futhark/TypeChecker/Terms.hs --- a/src/Language/Futhark/TypeChecker/Terms.hs +++ b/src/Language/Futhark/TypeChecker/Terms.hs @@ -37,7 +37,7 @@ import Language.Futhark.TypeChecker.Terms.Monad import Language.Futhark.TypeChecker.Terms.Pat import Language.Futhark.TypeChecker.Types -import Language.Futhark.TypeChecker.Unify hiding (Usage) +import Language.Futhark.TypeChecker.Unify import Prelude hiding (mod) overloadedTypeVars :: Constraints -> Names @@ -1331,8 +1331,12 @@ "Type is ambiguous (must be a sum type with constructors:" <+> ppr (Sum cs) <> ")." "Add a type annotation to disambiguate the type." - fixOverloaded (v, Size Nothing usage) = - typeError usage mempty $ "Size" <+> pquote (pprName v) <+> "is ambiguous.\n" + fixOverloaded (v, Size Nothing (Usage Nothing loc)) = + typeError loc mempty . withIndexLink "ambiguous-size" $ + "Ambiguous size" <+> pquote (pprName v) <> "." + fixOverloaded (v, Size Nothing (Usage (Just u) loc)) = + typeError loc mempty . withIndexLink "ambiguous-size" $ + "Ambiguous size" <+> pquote (pprName v) <+> "arising from" <+> text u <> "." fixOverloaded _ = pure () hiddenParamNames :: [Pat] -> Names diff --git a/src/Language/Futhark/TypeChecker/Terms/Monad.hs b/src/Language/Futhark/TypeChecker/Terms/Monad.hs --- a/src/Language/Futhark/TypeChecker/Terms/Monad.hs +++ b/src/Language/Futhark/TypeChecker/Terms/Monad.hs @@ -531,13 +531,14 @@ -- parameters. Returns the names of the fresh type variables, the -- instance list, and the instantiated type. instantiateTypeScheme :: + QualName VName -> SrcLoc -> [TypeParam] -> PatType -> TermTypeM ([VName], PatType) -instantiateTypeScheme loc tparams t = do +instantiateTypeScheme qn loc tparams t = do let tnames = map typeParamName tparams - (tparam_names, tparam_substs) <- unzip <$> mapM (instantiateTypeParam loc) tparams + (tparam_names, tparam_substs) <- unzip <$> mapM (instantiateTypeParam qn loc) tparams let substs = M.fromList $ zip tnames tparam_substs t' = applySubst (`M.lookup` substs) t pure (tparam_names, t') @@ -546,19 +547,22 @@ -- substitution map. instantiateTypeParam :: Monoid as => + QualName VName -> SrcLoc -> TypeParam -> TermTypeM (VName, Subst (RetTypeBase dim as)) -instantiateTypeParam loc tparam = do +instantiateTypeParam qn loc tparam = do i <- incCounter let name = nameFromString (takeWhile isAscii (baseString (typeParamName tparam))) v <- newID $ mkTypeVarName name i case tparam of TypeParamType x _ _ -> do - constrain v $ NoConstraint x $ mkUsage' loc + constrain v . NoConstraint x . mkUsage loc $ + "instantiated type parameter of " <> quote (pretty qn) <> "." pure (v, Subst [] $ RetType [] $ Scalar $ TypeVar mempty Nonunique (typeName v) []) TypeParamDim {} -> do - constrain v $ Size Nothing $ mkUsage' loc + constrain v . Size Nothing . mkUsage loc $ + "instantiated size parameter of " <> quote (pretty qn) <> "." pure (v, SizeSubst $ NamedDim $ qualName v) checkQualNameWithEnv :: Namespace -> QualName Name -> SrcLoc -> TermTypeM (TermScope, QualName VName) @@ -652,7 +656,7 @@ Just (BoundV _ tparams t) | "_" `isPrefixOf` baseString name -> underscoreUse loc qn | otherwise -> do - (tnames, t') <- instantiateTypeScheme loc tparams t + (tnames, t') <- instantiateTypeScheme qn' loc tparams t pure $ qualifyTypeVars outer_env tnames qs t' Just EqualityF -> do argtype <- newTypeVar loc "t" @@ -811,7 +815,7 @@ checkTypeExpNonrigid te = do (te', svars, RetType dims st) <- termCheckTypeExp te forM_ (svars ++ dims) $ \v -> - constrain v $ Size Nothing $ mkUsage' $ srclocOf te + constrain v $ Size Nothing $ mkUsage (srclocOf te) "anonymous size in type expression" pure (te', st, svars ++ dims) checkTypeExpRigid :: @@ -839,6 +843,8 @@ maybeDimFromExp e = ConstDim . fromIntegral <$> isInt64 e dimFromExp :: (Exp -> SizeSource) -> Exp -> TermTypeM (DimDecl VName, Maybe VName) +dimFromExp rf (Attr _ e _) = dimFromExp rf e +dimFromExp rf (Assert _ e _ _) = dimFromExp rf e dimFromExp rf (Parens e _) = dimFromExp rf e dimFromExp rf (QualParens _ e _) = dimFromExp rf e dimFromExp rf e diff --git a/src/Language/Futhark/TypeChecker/Terms/Pat.hs b/src/Language/Futhark/TypeChecker/Terms/Pat.hs --- a/src/Language/Futhark/TypeChecker/Terms/Pat.hs +++ b/src/Language/Futhark/TypeChecker/Terms/Pat.hs @@ -223,6 +223,9 @@ dimIdent _ NamedDim {} = Nothing patternDims _ = [] +sizeBinderToParam :: SizeBinder VName -> UncheckedTypeParam +sizeBinderToParam (SizeBinder v loc) = TypeParamDim (baseName v) loc + -- | Check and bind a @let@-pattern. bindingPat :: [SizeBinder VName] -> @@ -231,7 +234,7 @@ (Pat -> TermTypeM a) -> TermTypeM a bindingPat sizes p t m = do - checkForDuplicateNames [p] + checkForDuplicateNames (map sizeBinderToParam sizes) [p] checkPat sizes p t $ \p' -> binding (S.toList $ patIdents p') $ do -- Perform an observation of every declared dimension. This -- prevents unused-name warnings for otherwise unused dimensions. @@ -388,7 +391,7 @@ (Pat -> TermTypeM a) -> TermTypeM a checkPat sizes p t m = do - checkForDuplicateNames [p] + checkForDuplicateNames (map sizeBinderToParam sizes) [p] p' <- onFailure (CheckingPat p t) $ checkPat' sizes p t let explicit = mustBeExplicitInType $ patternStructType p' @@ -408,7 +411,7 @@ ([TypeParam] -> [Pat] -> TermTypeM a) -> TermTypeM a bindingParams tps orig_ps m = do - checkForDuplicateNames orig_ps + checkForDuplicateNames tps orig_ps checkTypeParams tps $ \tps' -> bindingTypeParams tps' $ do let descend ps' (p : ps) = checkPat [] p NoneInferred $ \p' -> diff --git a/src/Language/Futhark/TypeChecker/Types.hs b/src/Language/Futhark/TypeChecker/Types.hs --- a/src/Language/Futhark/TypeChecker/Types.hs +++ b/src/Language/Futhark/TypeChecker/Types.hs @@ -26,9 +26,10 @@ import Control.Monad.Reader import Control.Monad.State import Data.Bifunctor -import Data.List (foldl', sort, unzip4, (\\)) +import Data.List (find, foldl', sort, unzip4, (\\)) import qualified Data.Map.Strict as M import Data.Maybe +import qualified Data.Set as S import Futhark.Util (nubOrd) import Futhark.Util.Pretty hiding ((<|>)) import Language.Futhark @@ -233,12 +234,19 @@ dims' <- mapM (flip (checkName Term) loc) dims bindDims dims' $ do (t', svars, RetType t_dims st, l) <- evalTypeExp t - pure - ( TEDim dims' t' loc, - svars, - RetType (dims' ++ t_dims) st, - max l SizeLifted - ) + let (witnessed, _) = determineSizeWitnesses st + case find (`S.notMember` witnessed) dims' of + Just d -> + typeError loc mempty . withIndexLink "unused-existential" $ + "Existential size " <> pquote (pprName d) + <> " not used as array size." + Nothing -> + pure + ( TEDim dims' t' loc, + svars, + RetType (dims' ++ t_dims) st, + max l SizeLifted + ) where bindDims [] m = m bindDims (d : ds) m = @@ -341,26 +349,28 @@ checkForDuplicateNamesInType te evalTypeExp te --- | Check for duplication of names inside a pattern group. Produces --- a description of all names used in the pattern group. +-- | Check for duplication of names inside a binding group. checkForDuplicateNames :: - MonadTypeChecker m => - [UncheckedPat] -> - m () -checkForDuplicateNames = (`evalStateT` mempty) . mapM_ check + MonadTypeChecker m => [UncheckedTypeParam] -> [UncheckedPat] -> m () +checkForDuplicateNames tps pats = (`evalStateT` mempty) $ do + mapM_ checkTypeParam tps + mapM_ checkPat pats where - check (Id v _ loc) = seen v loc - check (PatParens p _) = check p - check (PatAttr _ p _) = check p - check Wildcard {} = pure () - check (TuplePat ps _) = mapM_ check ps - check (RecordPat fs _) = mapM_ (check . snd) fs - check (PatAscription p _ _) = check p - check PatLit {} = pure () - check (PatConstr _ _ ps _) = mapM_ check ps + checkTypeParam (TypeParamType _ v loc) = seen Type v loc + checkTypeParam (TypeParamDim v loc) = seen Term v loc - seen v loc = do - already <- gets $ M.lookup v + checkPat (Id v _ loc) = seen Term v loc + checkPat (PatParens p _) = checkPat p + checkPat (PatAttr _ p _) = checkPat p + checkPat Wildcard {} = pure () + checkPat (TuplePat ps _) = mapM_ checkPat ps + checkPat (RecordPat fs _) = mapM_ (checkPat . snd) fs + checkPat (PatAscription p _ _) = checkPat p + checkPat PatLit {} = pure () + checkPat (PatConstr _ _ ps _) = mapM_ checkPat ps + + seen ns v loc = do + already <- gets $ M.lookup (ns, v) case already of Just prev_loc -> lift $ @@ -368,7 +378,7 @@ "Name" <+> pquote (ppr v) <+> "also bound at" <+> text (locStr prev_loc) <> "." Nothing -> - modify $ M.insert v loc + modify $ M.insert (ns, v) loc -- | Check whether the type contains arrow types that define the same -- parameter. These might also exist further down, but that's not diff --git a/src/Language/Futhark/TypeChecker/Unify.hs b/src/Language/Futhark/TypeChecker/Unify.hs --- a/src/Language/Futhark/TypeChecker/Unify.hs +++ b/src/Language/Futhark/TypeChecker/Unify.hs @@ -8,7 +8,7 @@ -- blocks. module Language.Futhark.TypeChecker.Unify ( Constraint (..), - Usage, + Usage (..), mkUsage, mkUsage', Level, @@ -659,10 +659,18 @@ scopeCheck usage bcs vn lvl tp constraints <- getConstraints - let link = - let ext = filter (`S.member` typeDimNames tp) bound - in modifyConstraints $ + let link = do + let (witnessed, not_witnessed) = determineSizeWitnesses tp + used v = v `S.member` witnessed || v `S.member` not_witnessed + ext = filter used bound + case filter (`notElem` witnessed) ext of + [] -> + modifyConstraints $ M.insert vn (lvl, Constraint (RetType ext tp) usage) + problems -> + unifyError usage mempty bcs . withIndexLink "unify-param-existential" $ + "Parameter(s) " <> commasep (map (pquote . pprName) problems) + <> " used as size(s) would go out of scope." case snd <$> M.lookup vn constraints of Just (NoConstraint Unlifted unlift_usage) -> do @@ -1164,25 +1172,34 @@ indent 2 (ppr t2) "do not match." -runUnifyM :: [TypeParam] -> UnifyM a -> Either TypeError a -runUnifyM tparams (UnifyM m) = runExcept $ evalStateT m (constraints, 0) +runUnifyM :: [TypeParam] -> [TypeParam] -> UnifyM a -> Either TypeError a +runUnifyM rigid_tparams nonrigid_tparams (UnifyM m) = + runExcept $ evalStateT m (constraints, 0) where - constraints = M.fromList $ map f tparams - f (TypeParamDim p loc) = (p, (0, Size Nothing $ Usage Nothing loc)) - f (TypeParamType l p loc) = (p, (0, NoConstraint l $ Usage Nothing loc)) + constraints = + M.fromList $ + map nonrigid nonrigid_tparams <> map rigid rigid_tparams + nonrigid (TypeParamDim p loc) = (p, (0, Size Nothing $ Usage Nothing loc)) + nonrigid (TypeParamType l p loc) = (p, (0, NoConstraint l $ Usage Nothing loc)) + rigid (TypeParamDim p loc) = (p, (0, ParamSize loc)) + rigid (TypeParamType l p loc) = (p, (0, ParamType l loc)) -- | Perform a unification of two types outside a monadic context. --- The type parameters are allowed to be instantiated; all other types --- are considered rigid. +-- The first list of type parameters are rigid but may have liftedness +-- constraints; the second list of type parameters are allowed to be +-- instantiated. All other types are considered rigid with no +-- constraints. doUnification :: Loc -> [TypeParam] -> + [TypeParam] -> StructType -> StructType -> Either TypeError StructType -doUnification loc tparams t1 t2 = runUnifyM tparams $ do - expect (Usage Nothing (srclocOf loc)) t1 t2 - normTypeFully t2 +doUnification loc rigid_tparams nonrigid_tparams t1 t2 = + runUnifyM rigid_tparams nonrigid_tparams $ do + expect (Usage Nothing (srclocOf loc)) t1 t2 + normTypeFully t2 -- Note [Linking variables to sum types] -- diff --git a/src/Language/Futhark/Warnings.hs b/src/Language/Futhark/Warnings.hs --- a/src/Language/Futhark/Warnings.hs +++ b/src/Language/Futhark/Warnings.hs @@ -1,6 +1,6 @@ -- | A very simple representation of collections of warnings. -- Warnings have a position (so they can be ordered), and their --- 'Show'-instance produces a human-readable string. +-- 'Pretty'-instance produces a human-readable string. module Language.Futhark.Warnings ( Warnings, anyWarnings, diff --git a/src/main.hs b/src/main.hs --- a/src/main.hs +++ b/src/main.hs @@ -1,144 +1,8 @@ -{-# LANGUAGE OverloadedStrings #-} - --- | The @futhark@ command line program. +-- | The *actual* @futhark@ command line program, as seen by cabal. module 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 $ - [" 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 +import qualified Futhark.CLI.Main +-- | This is the main function. 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 +main = Futhark.CLI.Main.main diff --git a/unittests/Futhark/AD/DerivativesTests.hs b/unittests/Futhark/AD/DerivativesTests.hs new file mode 100644 --- /dev/null +++ b/unittests/Futhark/AD/DerivativesTests.hs @@ -0,0 +1,50 @@ +module Futhark.AD.DerivativesTests (tests) where + +import qualified Data.Map as M +import Futhark.AD.Derivatives +import Futhark.Analysis.PrimExp +import Futhark.IR.Syntax.Core (nameFromString) +import Futhark.Util.Pretty (pretty) +import Test.Tasty +import Test.Tasty.HUnit + +tests :: TestTree +tests = + testGroup + "Futhark.AD.DerivativesTests" + [ testGroup "Primitive functions" $ + map primFunTest $ + filter (not . (`elem` missing_primfuns) . fst) $ M.toList primFuns, + testGroup "BinOps" $ map binOpTest allBinOps, + testGroup "UnOps" $ map unOpTest allUnOps + ] + where + blank = ValueExp . blankPrimValue + + primFunTest (f, (ts, ret, _)) = + testCase f $ + case pdBuiltin (nameFromString f) (map blank ts) of + Nothing -> assertFailure "pdBuiltin gives Nothing" + Just v -> map primExpType v @?= replicate (length ts) ret + + -- We know we have no derivatives for these... and they are not + -- coming any time soon. + missing_primfuns = + [ "gamma16", + "gamma32", + "gamma64", + "lgamma16", + "lgamma32", + "lgamma64" + ] + + binOpTest bop = + testCase (pretty bop) $ + let t = binOpType bop + (dx, dy) = pdBinOp bop (blank t) (blank t) + in (primExpType dx, primExpType dy) @?= (t, t) + + unOpTest bop = + testCase (pretty bop) $ + let t = unOpType bop + in primExpType (pdUnOp bop $ blank t) @?= t diff --git a/unittests/Futhark/IR/Mem/IxFunTests.hs b/unittests/Futhark/IR/Mem/IxFunTests.hs --- a/unittests/Futhark/IR/Mem/IxFunTests.hs +++ b/unittests/Futhark/IR/Mem/IxFunTests.hs @@ -24,6 +24,7 @@ rem = P.rem div = P.div mod = P.mod + pow = (P.^) sgn = Just . P.signum allPoints :: [Int] -> [[Int]] diff --git a/unittests/Language/Futhark/TypeChecker/TypesTests.hs b/unittests/Language/Futhark/TypeChecker/TypesTests.hs --- a/unittests/Language/Futhark/TypeChecker/TypesTests.hs +++ b/unittests/Language/Futhark/TypeChecker/TypesTests.hs @@ -4,6 +4,7 @@ module Language.Futhark.TypeChecker.TypesTests (tests) where import Data.Bifunctor (first) +import Data.List (isInfixOf) import qualified Data.Map as M import Futhark.FreshNames import Futhark.Util.Pretty (prettyOneLine) @@ -16,13 +17,18 @@ import Test.Tasty import Test.Tasty.HUnit -evalTest :: TypeExp Name -> ([VName], StructRetType) -> TestTree +evalTest :: TypeExp Name -> Either String ([VName], StructRetType) -> TestTree evalTest te expected = testCase (pretty te) $ - case fmap (extract . fst) (run (checkTypeExp te)) of - Left e -> assertFailure $ "Failed: " <> pretty e - Right actual -> - actual @?= expected + case (fmap (extract . fst) (run (checkTypeExp te)), expected) of + (Left got_e, Left expected_e) -> + (expected_e `isInfixOf` pretty got_e) @? pretty got_e + (Left got_e, Right _) -> + assertFailure $ "Failed: " <> pretty got_e + (Right actual_t, Right expected_t) -> + actual_t @?= expected_t + (Right actual_t, Left _) -> + assertFailure $ "Expected error, got: " <> show actual_t where extract (_, svars, t, _) = (svars, t) run = snd . runTypeM env mempty (mkInitialImport "") blankNameSource @@ -68,67 +74,80 @@ evalTests = testGroup "Type expression elaboration" - [ evalTest - "[]i32" - ([], "?[d_0].[d_0]i32"), - evalTest - "[][]i32" - ([], "?[d_0][d_1].[d_0][d_1]i32"), - evalTest - "bool -> []i32" - ([], "bool -> ?[d_0].[d_0]i32"), - evalTest - "bool -> []f32 -> []i32" - (["d_0"], "bool -> [d_0]f32 -> ?[d_1].[d_1]i32"), - evalTest - "([]i32,[]i32)" - ([], "?[d_0][d_1].([d_0]i32, [d_1]i32)"), - evalTest - "{a:[]i32,b:[]i32}" - ([], "?[d_0][d_1].{a:[d_0]i32, b:[d_1]i32}"), - evalTest - "?[n].[n][n]bool" - ([], "?[n_0].[n_0][n_0]bool"), - evalTest - "([]i32 -> []i32) -> bool -> []i32" - (["d_0"], "([d_0]i32 -> ?[d_1].[d_1]i32) -> bool -> ?[d_2].[d_2]i32"), - evalTest - "((k: i64) -> [k]i32 -> [k]i32) -> []i32 -> bool" - (["d_1"], "((k_0: i64) -> [k_0]i32 -> [k_0]i32) -> [d_1]i32 -> bool"), - evalTest - "square [10]" - ([], "[10][10]i32"), - evalTest - "square []" - ([], "?[d_0].[d_0][d_0]i32"), - evalTest - "bool -> square []" - ([], "bool -> ?[d_0].[d_0][d_0]i32"), - evalTest - "(k: i64) -> square [k]" - ([], "(k_0: i64) -> [k_0][k_0]i32"), - evalTest - "fun i32 bool" - ([], "i32 -> bool"), - evalTest - "fun ([]i32) bool" - ([], "?[d_0].[d_0]i32 -> bool"), - evalTest - "fun bool ([]i32)" - ([], "?[d_0].bool -> [d_0]i32"), - evalTest - "bool -> fun ([]i32) bool" - ([], "bool -> ?[d_0].[d_0]i32 -> bool"), - evalTest - "bool -> fun bool ([]i32)" - ([], "bool -> ?[d_0].bool -> [d_0]i32"), - evalTest - "pair" - ([], "?[n_0][m_1].([n_0]i64, [m_1]i64)"), - evalTest - "(pair,pair)" - ([], "?[n_0][m_1][n_2][m_3].(([n_0]i64, [m_1]i64), ([n_2]i64, [m_3]i64))") + [ testGroup "Positive tests" (map mkPos pos), + testGroup "Negative tests" (map mkNeg neg) ] + where + mkPos (x, y) = evalTest x (Right y) + mkNeg (x, y) = evalTest x (Left y) + pos = + [ ( "[]i32", + ([], "?[d_0].[d_0]i32") + ), + ( "[][]i32", + ([], "?[d_0][d_1].[d_0][d_1]i32") + ), + ( "bool -> []i32", + ([], "bool -> ?[d_0].[d_0]i32") + ), + ( "bool -> []f32 -> []i32", + (["d_0"], "bool -> [d_0]f32 -> ?[d_1].[d_1]i32") + ), + ( "([]i32,[]i32)", + ([], "?[d_0][d_1].([d_0]i32, [d_1]i32)") + ), + ( "{a:[]i32,b:[]i32}", + ([], "?[d_0][d_1].{a:[d_0]i32, b:[d_1]i32}") + ), + ( "?[n].[n][n]bool", + ([], "?[n_0].[n_0][n_0]bool") + ), + ( "([]i32 -> []i32) -> bool -> []i32", + (["d_0"], "([d_0]i32 -> ?[d_1].[d_1]i32) -> bool -> ?[d_2].[d_2]i32") + ), + ( "((k: i64) -> [k]i32 -> [k]i32) -> []i32 -> bool", + (["d_1"], "((k_0: i64) -> [k_0]i32 -> [k_0]i32) -> [d_1]i32 -> bool") + ), + ( "square [10]", + ([], "[10][10]i32") + ), + ( "square []", + ([], "?[d_0].[d_0][d_0]i32") + ), + ( "bool -> square []", + ([], "bool -> ?[d_0].[d_0][d_0]i32") + ), + ( "(k: i64) -> square [k]", + ([], "(k_0: i64) -> [k_0][k_0]i32") + ), + ( "fun i32 bool", + ([], "i32 -> bool") + ), + ( "fun ([]i32) bool", + ([], "?[d_0].[d_0]i32 -> bool") + ), + ( "fun bool ([]i32)", + ([], "?[d_0].bool -> [d_0]i32") + ), + ( "bool -> fun ([]i32) bool", + ([], "bool -> ?[d_0].[d_0]i32 -> bool") + ), + ( "bool -> fun bool ([]i32)", + ([], "bool -> ?[d_0].bool -> [d_0]i32") + ), + ( "pair", + ([], "?[n_0][m_1].([n_0]i64, [m_1]i64)") + ), + ( "(pair,pair)", + ([], "?[n_0][m_1][n_2][m_3].(([n_0]i64, [m_1]i64), ([n_2]i64, [m_3]i64))") + ) + ] + neg = + [ ("?[n].bool", "Existential size \"n\""), + ("?[n].bool -> [n]bool", "Existential size \"n\""), + ("?[n].[n]bool -> [n]bool", "Existential size \"n\""), + ("?[n].[n]bool -> bool", "Existential size \"n\"") + ] substTest :: M.Map VName (Subst StructRetType) -> StructRetType -> StructRetType -> TestTree substTest m t expected = diff --git a/unittests/futhark_tests.hs b/unittests/futhark_tests.hs --- a/unittests/futhark_tests.hs +++ b/unittests/futhark_tests.hs @@ -1,5 +1,6 @@ module Main (main) where +import qualified Futhark.AD.DerivativesTests import qualified Futhark.BenchTests import qualified Futhark.IR.Mem.IxFunTests import qualified Futhark.IR.PrimitiveTests @@ -16,6 +17,7 @@ testGroup "" [ Language.Futhark.SyntaxTests.tests, + Futhark.AD.DerivativesTests.tests, Futhark.BenchTests.tests, Futhark.IR.PropTests.tests, Futhark.IR.Syntax.CoreTests.tests,