diff --git a/docs/language-reference.rst b/docs/language-reference.rst --- a/docs/language-reference.rst +++ b/docs/language-reference.rst @@ -554,8 +554,8 @@ `stringlit` ........... -Evaluates to an array of type ``[]i32`` that contains the code points -of the characters as integers. +Evaluates to an array of type ``[]u8`` that contains the characters +encoded as UTF-8. ``()`` ...... @@ -611,7 +611,7 @@ element. The ``s`` parameter may not be zero. If ``s`` is negative, it means to start at ``i`` and descend by steps of size ``s`` to ``j`` (not inclusive). Slicing can be done only with expressions of type -``i32``. +``i64``. It is generally a bad idea for ``s`` to be non-constant. Slicing of multiple dimensions can be done by separating with commas, @@ -782,7 +782,7 @@ Due to ambiguities, this syntactic form cannot appear as an array index expression unless it is first enclosed in parentheses. However, -as an array index must always be of type ``i32``, there is never a +as an array index must always be of type ``i64``, there is never a reason to put an explicit type ascription there. ``e :> t`` @@ -1033,6 +1033,9 @@ An application ``replicate 10 0`` will have type ``[10]i32``. +Whenever we write a type ``[n]t``, ``n`` must already be a variable of +type ``i64`` in scope (possibly by being bound as a size parameter). + .. _unknown-sizes: Unknown sizes @@ -1097,7 +1100,7 @@ Complex ranges .............. -Most complex ranges, such as ``a..` and :ref:`"upto" ranges `. diff --git a/docs/man/futhark-c.rst b/docs/man/futhark-c.rst --- a/docs/man/futhark-c.rst +++ b/docs/man/futhark-c.rst @@ -46,6 +46,9 @@ --safe Ignore ``unsafe`` in program and perform safety checks unconditionally. +--server + Generate a server-mode executable that reads commands from stdin. + -v verbose Enable debugging output. If compilation fails due to a compiler error, the result of the last successful compiler step will be diff --git a/docs/man/futhark-cuda.rst b/docs/man/futhark-cuda.rst --- a/docs/man/futhark-cuda.rst +++ b/docs/man/futhark-cuda.rst @@ -50,6 +50,9 @@ --safe Ignore ``unsafe`` in program and perform safety checks unconditionally. +--server + Generate a server-mode executable that reads commands from stdin. + -v verbose Enable debugging output. If compilation fails due to a compiler error, the result of the last successful compiler step will be diff --git a/docs/man/futhark-dataset.rst b/docs/man/futhark-dataset.rst --- a/docs/man/futhark-dataset.rst +++ b/docs/man/futhark-dataset.rst @@ -68,7 +68,7 @@ Generate an array of floating-point numbers and an array of indices into that array:: - futhark dataset -g [10]f32 --i32-bounds=0:9 -g [100]i32 + futhark dataset -g [10]f32 --i64-bounds=0:9 -g [100]i64 To generate binary data, the ``--binary`` must come before the ``--generate``:: diff --git a/docs/man/futhark-literate.rst b/docs/man/futhark-literate.rst --- a/docs/man/futhark-literate.rst +++ b/docs/man/futhark-literate.rst @@ -29,8 +29,11 @@ * Any *directives* will be executed and replaced with their output. See below. -**Warning:** Do not run untrusted programs. See SECURITY below. +**Warning:** Do not run untrusted programs. See SAFETY below. +Image directives and builtin functions Shells out to ``convert`` (from +ImageMagick). Video generation uses ``fmpeg``. + Directives ========== @@ -58,7 +61,7 @@ Shows the result of executing the FutharkScript expression ``e``, which can have any (transparent) type. -* ``> :anim e[; parameters...]`` +* ``> :video e[; parameters...]`` Creates a video from ``e``. The optional parameters are lines of the form *key: value*: @@ -69,8 +72,6 @@ * ``format: `` - Shells out to ``ffmpeg`` to actually create the video file. - ``e`` must be one of the following: * A 3D array where the 2D elements is of a type acceptable to @@ -94,8 +95,7 @@ * ``> :img e`` Visualises ``e``, which must be of type ``[][]i32`` or ``[][]u32`` - (interpreted as rows of ARGB pixel values). Shells out to - ``convert`` (from ImageMagick) to generate the image. + (interpreted as rows of ARGB pixel values). * ``> :plot2d e[; size=(height,width)]`` @@ -125,15 +125,22 @@ Only an extremely limited subset of Futhark is supported: .. productionlist:: - scriptexp: `id` `scriptexp`* + scriptexp: `fun` `scriptexp`* : | "(" `scriptexp` ")" : | "(" `scriptexp` ( "," `scriptexp` )+ ")" : | "{" "}" : | "{" (`id` = `scriptexp`) ("," `id` = `scriptexp`)* "}" : | `literal` + fun: `id` | "$" `id` -Any numeric literals *must* have a type suffix. +Function applications are either of Futhark funtions or *builtin +functions*. The latter are prefixed with ``$`` and are magical +(usually impure) functions that could not possibly be implemented in +Futhark. The following builtins are supported: +* ``$loadimg "file"`` reads an image from the given file and returns + it as a row-major ``[][]u32`` array with each pixel encoded as ARGB. + OPTIONS ======= @@ -180,12 +187,13 @@ Print verbose information on stderr about directives as they are executing. -SECURITY -======== +SAFETY +====== Some directives (e.g. ``:gnuplot``) can run arbitrary shell commands. -Running an untrusted literate Futhark program is as dangerous as -running a shell script you downloaded off the Internet. Before +Other directives or builtin functions can read or write arbitrary +files. Running an untrusted literate Futhark program is as dangerous +as running a shell script you downloaded off the Internet. Before running a program from an unknown source, you should always give it a quick read to see if anything looks fishy. diff --git a/docs/man/futhark-multicore.rst b/docs/man/futhark-multicore.rst --- a/docs/man/futhark-multicore.rst +++ b/docs/man/futhark-multicore.rst @@ -46,6 +46,9 @@ --safe Ignore ``unsafe`` in program and perform safety checks unconditionally. +--server + Generate a server-mode executable that reads commands from stdin. + -v verbose Enable debugging output. If compilation fails due to a compiler error, the result of the last successful compiler step will be diff --git a/docs/man/futhark-opencl.rst b/docs/man/futhark-opencl.rst --- a/docs/man/futhark-opencl.rst +++ b/docs/man/futhark-opencl.rst @@ -46,6 +46,9 @@ --safe Ignore ``unsafe`` in program and perform safety checks unconditionally. +--server + Generate a server-mode executable that reads commands from stdin. + -v verbose Enable debugging output. If compilation fails due to a compiler error, the result of the last successful compiler step will be diff --git a/docs/man/futhark-pyopencl.rst b/docs/man/futhark-pyopencl.rst --- a/docs/man/futhark-pyopencl.rst +++ b/docs/man/futhark-pyopencl.rst @@ -52,6 +52,9 @@ Where to write the resulting binary. By default, if the source program is named 'foo.fut', the binary will be named 'foo'. +--server + Generate a server-mode executable that reads commands from stdin. + --safe Ignore ``unsafe`` in program and perform safety checks unconditionally. diff --git a/docs/man/futhark-python.rst b/docs/man/futhark-python.rst --- a/docs/man/futhark-python.rst +++ b/docs/man/futhark-python.rst @@ -49,6 +49,9 @@ --safe Ignore ``unsafe`` in program and perform safety checks unconditionally. +--server + Generate a server-mode executable that reads commands from stdin. + -v verbose Enable debugging output. If compilation fails due to a compiler error, the result of the last successful compiler step will be 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 @@ -191,12 +191,12 @@ -- Test simple indexing of an array. -- == -- tags { firsttag secondtag } - -- input { [4,3,2,1] 1 } + -- input { [4,3,2,1] 1i64 } -- output { 3 } - -- input { [4,3,2,1] 5 } + -- input { [4,3,2,1] 5i64 } -- error: Assertion.*failed - let main (a: []i32) (i: i32): i32 = + let main (a: []i32) (i: i64): i32 = a[i] The following program contains two entry points, both of which are diff --git a/docs/server-protocol.rst b/docs/server-protocol.rst --- a/docs/server-protocol.rst +++ b/docs/server-protocol.rst @@ -23,12 +23,14 @@ Each command is sent as a *single line* on standard input. The response is sent on standard output. The server will print ``%%% OK`` -on a line by itself to indicate that a command has finished. If a -command fails, the server will print ``%%% FAILURE`` followed by the -error message, and then ``%%% OK`` when it is ready for more input. -Some output may also precede ``%%% FAILURE``, e.g. logging statements -that occured before failure was detected. Fatal errors (that lead to -server shutdown) may be printed to stderr. +on a line by itself to indicate that a command has finished. It will +also print ``%%% OK`` at startup once initialisation has finished. If +initialisation fails, the process will terminate. If a command fails, +the server will print ``%%% FAILURE`` followed by the error message, +and then ``%%% OK`` when it is ready for more input. Some output may +also precede ``%%% FAILURE``, e.g. logging statements that occured +before failure was detected. Fatal errors (that lead to server +shutdown) may be printed to stderr. Variables --------- diff --git a/futhark.cabal b/futhark.cabal --- a/futhark.cabal +++ b/futhark.cabal @@ -1,7 +1,7 @@ cabal-version: 2.4 - +-- Run 'cabal2nix . >futhark.nix' after adding deps. name: futhark -version: 0.18.6 +version: 0.19.1 synopsis: An optimising compiler for a functional, array-oriented language. description: Futhark is a small programming language designed to be compiled to @@ -106,6 +106,7 @@ Futhark.CodeGen.Backends.PyOpenCL Futhark.CodeGen.Backends.PyOpenCL.Boilerplate Futhark.CodeGen.Backends.SequentialC + Futhark.CodeGen.Backends.SequentialC.Boilerplate Futhark.CodeGen.Backends.SequentialPython Futhark.CodeGen.Backends.SimpleRep Futhark.CodeGen.ImpCode @@ -191,6 +192,7 @@ Futhark.Internalise.Monomorphise Futhark.Internalise.TypesValues Futhark.MonadFreshNames + Futhark.Optimise.BlkRegTiling Futhark.Optimise.CSE Futhark.Optimise.DoubleBuffer Futhark.Optimise.Fusion @@ -201,13 +203,18 @@ Futhark.Optimise.InPlaceLowering.SubstituteIndices Futhark.Optimise.InliningDeadFun Futhark.Optimise.Simplify - Futhark.Optimise.Simplify.ClosedForm Futhark.Optimise.Simplify.Engine Futhark.Optimise.Simplify.Lore Futhark.Optimise.Simplify.Rule Futhark.Optimise.Simplify.Rules + Futhark.Optimise.Simplify.Rules.BasicOp + Futhark.Optimise.Simplify.Rules.ClosedForm + Futhark.Optimise.Simplify.Rules.Index + Futhark.Optimise.Simplify.Rules.Loop + Futhark.Optimise.Simplify.Rules.Simple Futhark.Optimise.Sink Futhark.Optimise.TileLoops + Futhark.Optimise.TileLoops.Shared Futhark.Optimise.Unstream Futhark.Pass Futhark.Pass.ExpandAllocations @@ -294,6 +301,7 @@ , blaze-html >=0.9.0.1 , bytestring >=0.10.8 , bytestring-to-vector >=0.3.0.1 + , bmp >=1.2.6.3 , containers >=0.6.2.1 , directory >=1.3.0.0 , directory-tree >=0.12.1 diff --git a/prelude/array.fut b/prelude/array.fut --- a/prelude/array.fut +++ b/prelude/array.fut @@ -6,38 +6,62 @@ open import "zip" -- Rexport. -- | The size of the outer dimension of an array. +-- +-- **Complexity:** O(1). let length [n] 't (_: [n]t) = n -- | Is the array empty? +-- +-- **Complexity:** O(1). let null [n] 't (_: [n]t) = n == 0 -- | The first element of the array. +-- +-- **Complexity:** O(1). let head [n] 't (x: [n]t) = x[0] -- | The last element of the array. +-- +-- **Complexity:** O(1). let last [n] 't (x: [n]t) = x[n-1] -- | Everything but the first element of the array. +-- +-- **Complexity:** O(1). let tail [n] 't (x: [n]t) = x[1:] -- | Everything but the last element of the array. +-- +-- **Complexity:** O(1). let init [n] 't (x: [n]t) = x[0:n-1] -- | Take some number of elements from the head of the array. +-- +-- **Complexity:** O(1). let take [n] 't (i: i64) (x: [n]t): [i]t = x[0:i] -- | Remove some number of elements from the head of the array. +-- +-- **Complexity:** O(1). let drop [n] 't (i: i64) (x: [n]t) = x[i:] -- | Split an array at a given position. +-- +-- **Complexity:** O(1). let split [n] 't (i: i64) (xs: [n]t): ([i]t, []t) = (xs[:i] :> [i]t, xs[i:]) -- | Return the elements of the array in reverse order. +-- +-- **Complexity:** O(1). let reverse [n] 't (x: [n]t): [n]t = x[::-1] :> [n]t -- | Concatenate two arrays. Warning: never try to perform a reduction -- with this operator; it will not work. +-- +-- **Work:** O(n). +-- +-- **Span:** O(1). let (++) [n] [m] 't (xs: [n]t) (ys: [m]t): *[]t = intrinsics.concat (xs, ys) -- | An old-fashioned way of saying `++`. @@ -52,28 +76,48 @@ -- rotation amount is also supported. -- -- For example, if `b==rotate r a`, then `b[x+r] = a[x]`. +-- +-- **Complexity:** O(1). let rotate [n] 't (r: i64) (xs: [n]t): [n]t = intrinsics.rotate (r, xs) :> [n]t -- | Construct an array of consecutive integers of the given length, -- starting at 0. +-- +-- **Work:** O(n). +-- +-- **Span:** O(1). let iota (n: i64): *[n]i64 = 0..1.. [l]t --- | Combines the outer three dimensions of an array. +-- | Like `flatten`, but on the outer three dimensions of an array. let flatten_3d [n][m][l] 't (xs: [n][m][l]t): []t = flatten (flatten xs) --- | Combines the outer four dimensions of an array. +-- | Like `flatten`, but on the outer four dimensions of an array. let flatten_4d [n][m][l][k] 't (xs: [n][m][l][k]t): []t = flatten (flatten_3d xs) -- | Splits the outer dimension of an array in two. +-- +-- **Complexity:** O(1). let unflatten [p] 't (n: i64) (m: i64) (xs: [p]t): [n][m]t = intrinsics.unflatten (n, m, xs) :> [n][m]t --- | Splits the outer dimension of an array in three. +-- | Like `unflatten`, but produces three dimensions. let unflatten_3d [p] 't (n: i64) (m: i64) (l: i64) (xs: [p]t): [n][m][l]t = unflatten n m (unflatten (n*m) l xs) --- | Splits the outer dimension of an array in four. +-- | Like `unflatten`, but produces four dimensions. let unflatten_4d [p] 't (n: i64) (m: i64) (l: i64) (k: i64) (xs: [p]t): [n][m][l][k]t = unflatten n m (unflatten_3d (n*m) l k xs) +-- | Transpose an array. +-- +-- **Complexity:** O(1). let transpose [n] [m] 't (a: [n][m]t): [m][n]t = intrinsics.transpose a :> [m][n]t -- | True if all of the input elements are true. Produces true on an -- empty array. +-- +-- **Work:** O(n). +-- +-- **Span:** O(log(n)). let and [n] (xs: [n]bool) = all id xs -- | True if any of the input elements are true. Produces false on an -- empty array. +-- +-- **Work:** O(n). +-- +-- **Span:** O(log(n)). let or [n] (xs: [n]bool) = any id xs -- | Perform a *sequential* left-fold of an array. +-- +-- **Work:** O(n ✕ W(f))). +-- +-- **Span:** O(n ✕ S(f)). let foldl [n] 'a 'b (f: a -> b -> a) (acc: a) (bs: [n]b): a = loop acc for b in bs do f acc b -- | Perform a *sequential* right-fold of an array. +-- +-- **Work:** O(n ✕ W(f))). +-- +-- **Span:** O(n ✕ S(f)). let foldr [n] 'a 'b (f: b -> a -> a) (acc: a) (bs: [n]b): a = foldl (flip f) acc (reverse bs) -- | Create a value for each point in a one-dimensional index space. +-- +-- **Work:** *O(n ✕ W(f))* +-- +-- **Span:** *O(S(f))* let tabulate 'a (n: i64) (f: i64 -> a): *[n]a = map1 f (iota n) -- | Create a value for each point in a two-dimensional index space. +-- +-- **Work:** *O(n ✕ W(f))* +-- +-- **Span:** *O(S(f))* let tabulate_2d 'a (n: i64) (m: i64) (f: i64 -> i64 -> a): *[n][m]a = map1 (f >-> tabulate m) (iota n) -- | Create a value for each point in a three-dimensional index space. +-- +-- **Work:** *O(n ✕ W(f))* +-- +-- **Span:** *O(S(f))* let tabulate_3d 'a (n: i64) (m: i64) (o: i64) (f: i64 -> i64 -> i64 -> a): *[n][m][o]a = map1 (f >-> tabulate_2d m o) (iota n) diff --git a/prelude/math.fut b/prelude/math.fut --- a/prelude/math.fut +++ b/prelude/math.fut @@ -46,8 +46,6 @@ -- | Arithmetic negation (use `!` for bitwise negation). val neg: t -> t - -- | Deprecated alias for `neg`. - val negate: t -> t val max: t -> t -> t val min: t -> t -> t @@ -126,6 +124,9 @@ module type real = { include numeric + -- | Multiplicative inverse. + val recip: t -> t + val from_fraction: i64 -> i64 -> t val to_i64: t -> i64 val to_f64: t -> f64 @@ -282,7 +283,6 @@ let abs (x: i8) = intrinsics.abs8 x let neg (x: t) = -x - let negate = neg let max (x: t) (y: t) = intrinsics.smax8 (x, y) let min (x: t) (y: t) = intrinsics.smin8 (x, y) @@ -355,7 +355,6 @@ let abs (x: i16) = intrinsics.abs16 x let neg (x: t) = -x - let negate = neg let max (x: t) (y: t) = intrinsics.smax16 (x, y) let min (x: t) (y: t) = intrinsics.smin16 (x, y) @@ -431,7 +430,6 @@ let abs (x: i32) = intrinsics.abs32 x let neg (x: t) = -x - let negate = neg let max (x: t) (y: t) = intrinsics.smax32 (x, y) let min (x: t) (y: t) = intrinsics.smin32 (x, y) @@ -507,7 +505,6 @@ let abs (x: i64) = intrinsics.abs64 x let neg (x: t) = -x - let negate = neg let max (x: t) (y: t) = intrinsics.smax64 (x, y) let min (x: t) (y: t) = intrinsics.smin64 (x, y) @@ -583,7 +580,6 @@ let abs (x: u8) = x let neg (x: t) = -x - let negate = neg let max (x: t) (y: t) = unsign (intrinsics.umax8 (sign x, sign y)) let min (x: t) (y: t) = unsign (intrinsics.umin8 (sign x, sign y)) @@ -659,7 +655,6 @@ let abs (x: u16) = x let neg (x: t) = -x - let negate = neg let max (x: t) (y: t) = unsign (intrinsics.umax16 (sign x, sign y)) let min (x: t) (y: t) = unsign (intrinsics.umin16 (sign x, sign y)) @@ -738,7 +733,6 @@ let lowest = highest + 1u32 let neg (x: t) = -x - let negate = neg let max (x: t) (y: t) = unsign (intrinsics.umax32 (sign x, sign y)) let min (x: t) (y: t) = unsign (intrinsics.umin32 (sign x, sign y)) @@ -811,7 +805,6 @@ let abs (x: u64) = x let neg (x: t) = -x - let negate = neg let max (x: t) (y: t) = unsign (intrinsics.umax64 (sign x, sign y)) let min (x: t) (y: t) = unsign (intrinsics.umin64 (sign x, sign y)) @@ -875,13 +868,11 @@ let (x: f64) != (y: f64) = intrinsics.! (x == y) let neg (x: t) = -x - let negate = neg + let recip (x: t) = 1/x let max (x: t) (y: t) = intrinsics.fmax64 (x, y) let min (x: t) (y: t) = intrinsics.fmin64 (x, y) - let sgn (x: f64) = if x < 0f64 then -1f64 - else if x == 0f64 then 0f64 - else 1f64 + let sgn (x: f64) = intrinsics.fsignum64 x let abs (x: f64) = intrinsics.fabs64 x let sqrt (x: f64) = intrinsics.sqrt64 x @@ -983,13 +974,11 @@ let (x: f32) != (y: f32) = intrinsics.! (x == y) let neg (x: t) = -x - let negate = neg + let recip (x: t) = 1/x let max (x: t) (y: t) = intrinsics.fmax32 (x, y) let min (x: t) (y: t) = intrinsics.fmin32 (x, y) - let sgn (x: f32) = if x < 0f32 then -1f32 - else if x == 0f32 then 0f32 - else 1f32 + let sgn (x: f32) = intrinsics.fsignum32 x let abs (x: f32) = intrinsics.fabs32 x let sqrt (x: f32) = intrinsics.sqrt32 x diff --git a/prelude/soacs.fut b/prelude/soacs.fut --- a/prelude/soacs.fut +++ b/prelude/soacs.fut @@ -9,9 +9,20 @@ -- slower than `reduce`@term, although they have the same asymptotic -- complexity. -- --- *Reminder on terminology*: A function `op` is said to be --- *associative* if +-- **Higher-order complexity** -- +-- Specifying the time complexity of higher-order functions is tricky +-- because it depends on the functional argument. We use the informal +-- convention that *W(f)* denotes the largest (asymptotic) *work* of +-- function *f*, for the values it may be applied to. Similarly, +-- *S(f)* denotes the largest span. See [this Wikipedia +-- article](https://en.wikipedia.org/wiki/Analysis_of_parallel_algorithms) +-- for a general introduction to these constructs. +-- +-- **Reminder on terminology** +-- +-- A function `op` is said to be *associative* if +-- -- (x `op` y) `op` z == x `op` (y `op` z) -- -- for all `x`, `y`, `z`. Similarly, it is *commutative* if @@ -33,49 +44,49 @@ -- | Apply the given function to each element of an array. -- --- **Work:** *O(n)* +-- **Work:** *O(n ✕ W(f))* -- --- **Span:** *O(1)* +-- **Span:** *O(S(f))* let map 'a [n] 'x (f: a -> x) (as: [n]a): *[n]x = intrinsics.map (f, as) :> *[n]x -- | Apply the given function to each element of a single array. -- --- **Work:** *O(n)* +-- **Work:** *O(n ✕ W(f))* -- --- **Span:** *O(1)* +-- **Span:** *O(S(f))* let map1 'a [n] 'x (f: a -> x) (as: [n]a): *[n]x = map f as -- | As `map1`@term, but with one more array. -- --- **Work:** *O(n)* +-- **Work:** *O(n ✕ W(f))* -- --- **Span:** *O(1)* +-- **Span:** *O(S(f))* let map2 'a 'b [n] 'x (f: a -> b -> x) (as: [n]a) (bs: [n]b): *[n]x = map (\(a, b) -> f a b) (zip2 as bs) -- | As `map2`@term, but with one more array. -- --- **Work:** *O(n)* +-- **Work:** *O(n ✕ W(f))* -- --- **Span:** *O(1)* +-- **Span:** *O(S(f))* let map3 'a 'b 'c [n] 'x (f: a -> b -> c -> x) (as: [n]a) (bs: [n]b) (cs: [n]c): *[n]x = map (\(a, b, c) -> f a b c) (zip3 as bs cs) -- | As `map3`@term, but with one more array. -- --- **Work:** *O(n)* +-- **Work:** *O(n ✕ W(f))* -- --- **Span:** *O(1)* +-- **Span:** *O(S(f))* let map4 'a 'b 'c 'd [n] 'x (f: a -> b -> c -> d -> x) (as: [n]a) (bs: [n]b) (cs: [n]c) (ds: [n]d): *[n]x = map (\(a, b, c, d) -> f a b c d) (zip4 as bs cs ds) --- | As `map4`@term, but with one more array. +-- | As `map3`@term, but with one more array. -- --- **Work:** *O(n)* +-- **Work:** *O(n ✕ W(f))* -- --- **Span:** *O(1)* +-- **Span:** *O(S(f))* let map5 'a 'b 'c 'd 'e [n] 'x (f: a -> b -> c -> d -> e -> x) (as: [n]a) (bs: [n]b) (cs: [n]c) (ds: [n]d) (es: [n]e): *[n]x = map (\(a, b, c, d, e) -> f a b c d e) (zip5 as bs cs ds es) @@ -86,9 +97,12 @@ -- the neutral element, and that must again have the same size as the -- elements of the input array. -- --- **Work:** *O(n)* +-- **Work:** *O(n ✕ W(op))* -- --- **Span:** *O(log(n))* +-- **Span:** *O(log(n) ✕ W(op))* +-- +-- Note that the complexity implies that parallelism in the combining +-- operator will *not* be exploited. let reduce [n] 'a (op: a -> a -> a) (ne: a) (as: [n]a): a = intrinsics.reduce (op, ne, as) @@ -97,9 +111,9 @@ -- like addition, the compiler already knows that the operator is -- commutative, so plain `reduce`@term will work just as well. -- --- **Work:** *O(n)* +-- **Work:** *O(n ✕ W(op))* -- --- **Span:** *O(log(n))* +-- **Span:** *O(log(n) ✕ W(op))* let reduce_comm [n] 'a (op: a -> a -> a) (ne: a) (as: [n]a): a = intrinsics.reduce_comm (op, ne, as) @@ -111,10 +125,10 @@ -- associative and commutative. Out-of-bounds indices in `is` are -- ignored. -- --- **Work:** *O(n)* +-- **Work:** *O(n ✕ W(op))* -- --- **Span:** *O(n)* in the worst case (all updates to same position), --- but *O(1)* in the best case. +-- **Span:** *O(n ✕ W(op))* in the worst case (all updates to same +-- position), but *O(W(op))* in the best case. -- -- In practice, the *O(n)* behaviour only occurs if *m* is also very -- large. @@ -122,20 +136,20 @@ intrinsics.hist (1, dest, f, ne, is, as) :> *[m]a -- | Inclusive prefix scan. Has the same caveats with respect to --- associativity as `reduce`. +-- associativity and complexity as `reduce`. -- --- **Work:** *O(n)* +-- **Work:** *O(n ✕ W(op))* -- --- **Span:** *O(log(n))* +-- **Span:** *O(log(n) ✕ W(op))* let scan [n] 'a (op: a -> a -> a) (ne: a) (as: [n]a): *[n]a = intrinsics.scan (op, ne, as) :> *[n]a -- | Remove all those elements of `as` that do not satisfy the -- predicate `p`. -- --- **Work:** *O(n)* +-- **Work:** *O(n ✕ W(p))* -- --- **Span:** *O(log(n))* +-- **Span:** *O(log(n) ✕ W(p))* let filter [n] 'a (p: a -> bool) (as: [n]a): *[]a = let (as', is) = intrinsics.partition (1, \x -> if p x then 0 else 1, as) in as'[:is[0]] @@ -143,9 +157,9 @@ -- | Split an array into those elements that satisfy the given -- predicate, and those that do not. -- --- **Work:** *O(n)* +-- **Work:** *O(n ✕ W(p))* -- --- **Span:** *O(log(n))* +-- **Span:** *O(log(n) ✕ W(p))* let partition [n] 'a (p: a -> bool) (as: [n]a): ([]a, []a) = let p' x = if p x then 0 else 1 let (as', is) = intrinsics.partition (2, p', as) @@ -153,9 +167,9 @@ -- | Split an array by two predicates, producing three arrays. -- --- **Work:** *O(n)* +-- **Work:** *O(n ✕ (W(p1) + W(p2)))* -- --- **Span:** *O(log(n))* +-- **Span:** *O(log(n) ✕ (W(p1) + W(p2)))* let partition2 [n] 'a (p1: a -> bool) (p2: a -> bool) (as: [n]a): ([]a, []a, []a) = let p' x = if p1 x then 0 else if p2 x then 1 else 2 let (as', is) = intrinsics.partition (3, p', as) @@ -173,9 +187,9 @@ -- A chunk may be empty, and `f 0 []` must produce the neutral element for -- `op`. -- --- **Work:** *O(n)* +-- **Work:** *O(n ✕ W(op) + W(f))* -- --- **Span:** *O(log(n))* +-- **Span:** *O(log(n) ✕ W(op))* let reduce_stream [n] 'a 'b (op: b -> b -> b) (f: (k: i64) -> [k]a -> b) (as: [n]a): b = intrinsics.reduce_stream (op, f, as) @@ -183,9 +197,9 @@ -- correspond to subsequences of the original array (they may be -- interleaved). -- --- **Work:** *O(n)* +-- **Work:** *O(n ✕ W(op) + W(f))* -- --- **Span:** *O(log(n))* +-- **Span:** *O(log(n) ✕ W(op))* let reduce_stream_per [n] 'a 'b (op: b -> b -> b) (f: (k: i64) -> [k]a -> b) (as: [n]a): b = intrinsics.reduce_stream_per (op, f, as) @@ -193,9 +207,9 @@ -- an array *of the same size*. The per-chunk results are -- concatenated. -- --- **Work:** *O(n)* +-- **Work:** *O(n ✕ W(f))* -- --- **Span:** *O(1)* +-- **Span:** *O(S(f))* let map_stream [n] 'a 'b (f: (k: i64) -> [k]a -> [k]b) (as: [n]a): *[n]b = intrinsics.map_stream (f, as) :> *[n]b @@ -203,32 +217,32 @@ -- correspond to subsequences of the original array (they may be -- interleaved). -- --- **Work:** *O(n)* +-- **Work:** *O(n ✕ W(f))* -- --- **Span:** *O(1)* +-- **Span:** *O(S(f))* let 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 -- | Return `true` if the given function returns `true` for all -- elements in the array. -- --- **Work:** *O(n)* +-- **Work:** *O(n ✕ W(f))* -- --- **Span:** *O(log(n))* +-- **Span:** *O(log(n) + S(f))* let all [n] 'a (f: a -> bool) (as: [n]a): bool = reduce (&&) true (map f as) -- | Return `true` if the given function returns `true` for any -- elements in the array. -- --- **Work:** *O(n)* +-- **Work:** *O(n ✕ W(f))* -- --- **Span:** *O(log(n))* +-- **Span:** *O(log(n) + S(f))* let any [n] 'a (f: a -> bool) (as: [n]a): bool = reduce (||) false (map f as) --- | The `scatter as is vs` expression calculates the equivalent of --- this imperative code: +-- | `scatter as is vs` calculates the equivalent of this imperative +-- code: -- -- ``` -- for index in 0..length is-1: @@ -254,3 +268,21 @@ -- **Span:** *O(1)* let scatter 't [m] [n] (dest: *[m]t) (is: [n]i64) (vs: [n]t): *[m]t = intrinsics.scatter (dest, is, vs) :> *[m]t + +-- | `scatter_2d as is vs` is the equivalent of a `scatter` on a 2-dimensional +-- array. +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(1)* +let 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 + +-- | `scatter_3d as is vs` is the equivalent of a `scatter` on a 3-dimensional +-- array. +-- +-- **Work:** *O(n)* +-- +-- **Span:** *O(1)* +let 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 diff --git a/prelude/zip.fut b/prelude/zip.fut --- a/prelude/zip.fut +++ b/prelude/zip.fut @@ -1,6 +1,8 @@ -- | Transforming arrays of tuples into tuples of arrays and back --- again. These are generally very cheap operations, as the internal --- compiler representation is always tuples of arrays. +-- again. +-- +-- These are generally very cheap operations, as the internal compiler +-- representation is always tuples of arrays. -- The main reason this module exists is that we need it to define -- SOACs like `map2`. diff --git a/rts/c/atomics.h b/rts/c/atomics.h --- a/rts/c/atomics.h +++ b/rts/c/atomics.h @@ -1,5 +1,39 @@ // Start of atomics.h +inline int32_t atomic_xchg_i32_global(volatile __global int32_t *p, int32_t x) { +#ifdef FUTHARK_CUDA + return atomicExch((int32_t*)p, x); +#else + return atomic_xor(p, x); +#endif +} + +inline int32_t atomic_xchg_i32_local(volatile __local int32_t *p, int32_t x) { +#ifdef FUTHARK_CUDA + return atomicExch((int32_t*)p, x); +#else + return atomic_xor(p, x); +#endif +} + +inline int32_t atomic_cmpxchg_i32_global(volatile __global int32_t *p, + int32_t cmp, int32_t val) { +#ifdef FUTHARK_CUDA + return atomicCAS((int32_t*)p, cmp, val); +#else + return atomic_cmpxchg(p, cmp, val); +#endif +} + +inline int32_t atomic_cmpxchg_i32_local(volatile __local int32_t *p, + int32_t cmp, int32_t val) { +#ifdef FUTHARK_CUDA + return atomicCAS((int32_t*)p, cmp, val); +#else + return atomic_cmpxchg(p, cmp, val); +#endif +} + inline int32_t atomic_add_i32_global(volatile __global int32_t *p, int32_t x) { #ifdef FUTHARK_CUDA return atomicAdd((int32_t*)p, x); @@ -26,7 +60,7 @@ do { assumed.f = old.f; old.f = old.f + x; - old.i = atomic_cmpxchg((volatile __global int32_t*)p, assumed.i, old.i); + old.i = atomic_cmpxchg_i32_global((volatile __global int32_t*)p, assumed.i, old.i); } while (assumed.i != old.i); return old.f; #endif @@ -42,7 +76,7 @@ do { assumed.f = old.f; old.f = old.f + x; - old.i = atomic_cmpxchg((volatile __local int32_t*)p, assumed.i, old.i); + old.i = atomic_cmpxchg_i32_local((volatile __local int32_t*)p, assumed.i, old.i); } while (assumed.i != old.i); return old.f; #endif @@ -160,37 +194,199 @@ #endif } -inline int32_t atomic_xchg_i32_global(volatile __global int32_t *p, int32_t x) { +// Start of 64 bit atomics + +inline int64_t atomic_xchg_i64_global(volatile __global int64_t *p, int64_t x) { #ifdef FUTHARK_CUDA - return atomicExch((int32_t*)p, x); + return atomicExch((uint64_t*)p, x); #else - return atomic_xor(p, x); + return atom_xor(p, x); #endif } -inline int32_t atomic_xchg_i32_local(volatile __local int32_t *p, int32_t x) { +inline int64_t atomic_xchg_i64_local(volatile __local int64_t *p, int64_t x) { #ifdef FUTHARK_CUDA - return atomicExch((int32_t*)p, x); + return atomicExch((uint64_t*)p, x); #else - return atomic_xor(p, x); + return atom_xor(p, x); #endif } -inline int32_t atomic_cmpxchg_i32_global(volatile __global int32_t *p, - int32_t cmp, int32_t val) { +inline int64_t atomic_cmpxchg_i64_global(volatile __global int64_t *p, + int64_t cmp, int64_t val) { #ifdef FUTHARK_CUDA - return atomicCAS((int32_t*)p, cmp, val); + return atomicCAS((uint64_t*)p, cmp, val); #else - return atomic_cmpxchg(p, cmp, val); + return atom_cmpxchg(p, cmp, val); #endif } -inline int32_t atomic_cmpxchg_i32_local(volatile __local int32_t *p, - int32_t cmp, int32_t val) { +inline int64_t atomic_cmpxchg_i64_local(volatile __local int64_t *p, + int64_t cmp, int64_t val) { #ifdef FUTHARK_CUDA - return atomicCAS((int32_t*)p, cmp, val); + return atomicCAS((uint64_t*)p, cmp, val); #else - return atomic_cmpxchg(p, cmp, val); + return atom_cmpxchg(p, cmp, val); +#endif +} + +inline int64_t atomic_add_i64_global(volatile __global int64_t *p, int64_t x) { +#ifdef FUTHARK_CUDA + return atomicAdd((uint64_t*)p, x); +#else + return atom_add(p, x); +#endif +} + +inline int64_t atomic_add_i64_local(volatile __local int64_t *p, int64_t x) { +#ifdef FUTHARK_CUDA + return atomicAdd((uint64_t*)p, x); +#else + return atom_add(p, x); +#endif +} + +inline double atomic_fadd_f64_global(volatile __global double *p, double x) { +#if defined(FUTHARK_CUDA) && __CUDA_ARCH__ >= 600 + return atomicAdd((double*)p, x); +#else + union { int64_t i; double f; } old; + union { int64_t i; double f; } assumed; + old.f = *p; + do { + assumed.f = old.f; + old.f = old.f + x; + old.i = atomic_cmpxchg_i64_global((volatile __global int64_t*)p, assumed.i, old.i); + } while (assumed.i != old.i); + return old.f; +#endif +} + +inline double atomic_fadd_f64_local(volatile __local double *p, double x) { +#if defined(FUTHARK_CUDA) && __CUDA_ARCH__ >= 600 + return atomicAdd((double*)p, x); +#else + union { int64_t i; double f; } old; + union { int64_t i; double f; } assumed; + old.f = *p; + do { + assumed.f = old.f; + old.f = old.f + x; + old.i = atomic_cmpxchg_i64_local((volatile __local int64_t*)p, assumed.i, old.i); + } while (assumed.i != old.i); + return old.f; +#endif +} + +inline int64_t atomic_smax_i64_global(volatile __global int64_t *p, int64_t x) { +#ifdef FUTHARK_CUDA + return atomicMax((int64_t*)p, x); +#else + return atom_max(p, x); +#endif +} + +inline int64_t atomic_smax_i64_local(volatile __local int64_t *p, int64_t x) { +#ifdef FUTHARK_CUDA + return atomicMax((int64_t*)p, x); +#else + return atom_max(p, x); +#endif +} + +inline int64_t atomic_smin_i64_global(volatile __global int64_t *p, int64_t x) { +#ifdef FUTHARK_CUDA + return atomicMin((int64_t*)p, x); +#else + return atom_min(p, x); +#endif +} + +inline int64_t atomic_smin_i64_local(volatile __local int64_t *p, int64_t x) { +#ifdef FUTHARK_CUDA + return atomicMin((int64_t*)p, x); +#else + return atom_min(p, x); +#endif +} + +inline uint64_t atomic_umax_i64_global(volatile __global uint64_t *p, uint64_t x) { +#ifdef FUTHARK_CUDA + return atomicMax((uint64_t*)p, x); +#else + return atom_max(p, x); +#endif +} + +inline uint64_t atomic_umax_i64_local(volatile __local uint64_t *p, uint64_t x) { +#ifdef FUTHARK_CUDA + return atomicMax((uint64_t*)p, x); +#else + return atom_max(p, x); +#endif +} + +inline uint64_t atomic_umin_i64_global(volatile __global uint64_t *p, uint64_t x) { +#ifdef FUTHARK_CUDA + return atomicMin((uint64_t*)p, x); +#else + return atom_min(p, x); +#endif +} + +inline uint64_t atomic_umin_i64_local(volatile __local uint64_t *p, uint64_t x) { +#ifdef FUTHARK_CUDA + return atomicMin((uint64_t*)p, x); +#else + return atom_min(p, x); +#endif +} + +inline int64_t atomic_and_i64_global(volatile __global int64_t *p, int64_t x) { +#ifdef FUTHARK_CUDA + return atomicAnd((int64_t*)p, x); +#else + return atom_and(p, x); +#endif +} + +inline int64_t atomic_and_i64_local(volatile __local int64_t *p, int64_t x) { +#ifdef FUTHARK_CUDA + return atomicAnd((int64_t*)p, x); +#else + return atom_and(p, x); +#endif +} + +inline int64_t atomic_or_i64_global(volatile __global int64_t *p, int64_t x) { +#ifdef FUTHARK_CUDA + return atomicOr((int64_t*)p, x); +#else + return atom_or(p, x); +#endif +} + +inline int64_t atomic_or_i64_local(volatile __local int64_t *p, int64_t x) { +#ifdef FUTHARK_CUDA + return atomicOr((int64_t*)p, x); +#else + return atom_or(p, x); +#endif +} + +inline int64_t atomic_xor_i64_global(volatile __global int64_t *p, int64_t x) { +#ifdef FUTHARK_CUDA + return atomicXor((int64_t*)p, x); +#else + return atom_xor(p, x); +#endif +} + +inline int64_t atomic_xor_i64_local(volatile __local int64_t *p, int64_t x) { +#ifdef FUTHARK_CUDA + return atomicXor((int64_t*)p, x); +#else + return atom_xor(p, x); #endif } diff --git a/rts/c/cuda.h b/rts/c/cuda.h --- a/rts/c/cuda.h +++ b/rts/c/cuda.h @@ -38,6 +38,7 @@ size_t default_block_size; size_t default_grid_size; size_t default_tile_size; + size_t default_reg_tile_size; size_t default_threshold; int default_block_size_changed; @@ -70,6 +71,7 @@ cfg->default_block_size = 256; cfg->default_grid_size = 0; // Set properly later. cfg->default_tile_size = 32; + cfg->default_reg_tile_size = 2; cfg->default_threshold = 32*1024; cfg->default_block_size_changed = 0; @@ -408,6 +410,9 @@ } else if (strstr(size_class, "tile_size") == size_class) { max_value = ctx->max_tile_size; default_value = ctx->cfg.default_tile_size; + } else if (strstr(size_class, "reg_tile_size") == size_class) { + max_value = 0; // No limit. + default_value = ctx->cfg.default_reg_tile_size; } else if (strstr(size_class, "threshold") == size_class) { // Threshold can be as large as it takes. default_value = ctx->cfg.default_threshold; diff --git a/rts/c/opencl.h b/rts/c/opencl.h --- a/rts/c/opencl.h +++ b/rts/c/opencl.h @@ -38,6 +38,7 @@ size_t default_group_size; size_t default_num_groups; size_t default_tile_size; + size_t default_reg_tile_size; size_t default_threshold; int default_group_size_changed; @@ -74,6 +75,7 @@ cfg->default_group_size = 0; cfg->default_num_groups = 0; cfg->default_tile_size = 0; + cfg->default_reg_tile_size = 0; cfg->default_threshold = 0; cfg->default_group_size_changed = 0; @@ -617,6 +619,9 @@ } else if (strstr(size_class, "tile_size") == size_class) { max_value = sqrt(max_group_size); default_value = ctx->cfg.default_tile_size; + } else if (strstr(size_class, "reg_tile_size") == size_class) { + max_value = 0; // No limit. + default_value = ctx->cfg.default_reg_tile_size; } else if (strstr(size_class, "threshold") == size_class) { // Threshold can be as large as it takes. default_value = ctx->cfg.default_threshold; diff --git a/rts/c/server.h b/rts/c/server.h --- a/rts/c/server.h +++ b/rts/c/server.h @@ -531,7 +531,6 @@ } else { futhark_panic(1, "Unknown command: %s\n", command); } - ok(); } void run_server(struct futhark_prog *prog, struct futhark_context *ctx) { @@ -551,8 +550,10 @@ s.variables[i].name = NULL; } + ok(); while ((linelen = getline(&line, &buflen, stdin)) > 0) { process_line(&s, line); + ok(); } free(line); diff --git a/rts/c/tuning.h b/rts/c/tuning.h --- a/rts/c/tuning.h +++ b/rts/c/tuning.h @@ -21,9 +21,10 @@ *eql = 0; int value = atoi(eql+1); if (set_size(cfg, line, value) != 0) { - strncpy(eql+1, line, max_line_len-strlen(line)-1); - snprintf(line, max_line_len, "Unknown name '%s' on line %d.", eql+1, lineno); - return line; + char* err = (char*) malloc(max_line_len + 50); + snprintf(err, max_line_len + 50, "Unknown name '%s' on line %d.", line, lineno); + free(line); + return err; } } else { snprintf(line, max_line_len, "Invalid line %d (must be of form 'name=int').", diff --git a/rts/python/opencl.py b/rts/python/opencl.py --- a/rts/python/opencl.py +++ b/rts/python/opencl.py @@ -81,6 +81,7 @@ default_group_size=None, default_num_groups=None, default_tile_size=None, + default_reg_tile_size=None, default_threshold=None, size_heuristics=[], required_types=[], @@ -138,6 +139,10 @@ default_tile_size = sizes['default_tile_size'] del sizes['default_tile_size'] + if 'default_reg_tile_size' in sizes: + default_reg_tile_size = sizes['default_reg_tile_size'] + del sizes['default_reg_tile_size'] + if 'default_threshold' in sizes: default_threshold = sizes['default_threshold'] del sizes['default_threshold'] @@ -147,6 +152,7 @@ default_sizes = apply_size_heuristics(self, size_heuristics, {'group_size': default_group_size, 'tile_size': default_tile_size, + 'reg_tile_size': default_reg_tile_size, 'num_groups': default_num_groups, 'lockstep_width': None, 'threshold': default_threshold}) @@ -154,6 +160,7 @@ default_num_groups = default_sizes['num_groups'] default_threshold = default_sizes['threshold'] default_tile_size = default_sizes['tile_size'] + default_reg_tile_size = default_sizes['reg_tile_size'] lockstep_width = default_sizes['lockstep_width'] if default_group_size > max_group_size: @@ -185,6 +192,9 @@ elif v['class'] == 'tile_size': max_value = max_tile_size default_value = default_tile_size + elif v['class'] == 'reg_tile_size': + max_value = None + default_value = default_reg_tile_size elif v['class'].startswith('threshold'): max_value = None default_value = default_threshold diff --git a/rts/python/server.py b/rts/python/server.py --- a/rts/python/server.py +++ b/rts/python/server.py @@ -151,6 +151,7 @@ def run(self): while True: + print('%%% OK', flush=True) line = sys.stdin.readline() if line == '': return @@ -159,7 +160,5 @@ except self.Failure as e: print('%%% FAILURE') print(e.msg) - print('%%% OK', flush=True) - # End of server.py diff --git a/src/Futhark/Analysis/Alias.hs b/src/Futhark/Analysis/Alias.hs --- a/src/Futhark/Analysis/Alias.hs +++ b/src/Futhark/Analysis/Alias.hs @@ -12,7 +12,6 @@ ( aliasAnalysis, -- * Ad-hoc utilities - AliasTable, analyseFun, analyseStms, analyseExp, @@ -42,10 +41,6 @@ where body' = analyseBody mempty body --- | Pre-existing aliases for variables. Used to add transitive --- aliases. -type AliasTable = M.Map VName Names - analyseBody :: ( ASTLore lore, CanBeAliased (Op lore) @@ -116,18 +111,16 @@ mapOnBranchType = return, mapOnFParam = return, mapOnLParam = return, - mapOnOp = return . addOpAliases + mapOnOp = return . addOpAliases aliases } analyseLambda :: (ASTLore lore, CanBeAliased (Op lore)) => + AliasTable -> Lambda lore -> Lambda (Aliases lore) -analyseLambda lam = - -- XXX: it may cause trouble that we pass mempty to analyseBody - -- here. However, fixing this generally involves adding an - -- AliasTable argument to addOpAliases. - let body = analyseBody mempty $ lambdaBody lam +analyseLambda aliases lam = + let body = analyseBody aliases $ lambdaBody lam in lam { lambdaBody = body, lambdaParams = lambdaParams lam diff --git a/src/Futhark/Analysis/HORep/SOAC.hs b/src/Futhark/Analysis/HORep/SOAC.hs --- a/src/Futhark/Analysis/HORep/SOAC.hs +++ b/src/Futhark/Analysis/HORep/SOAC.hs @@ -87,7 +87,6 @@ ScremaForm (..), StreamForm (..), StreamOrd (..), - getStreamAccums, scremaType, ) import qualified Futhark.IR.SOACS.SOAC as Futhark @@ -376,8 +375,8 @@ -- | A definite representation of a SOAC expression. data SOAC lore - = Stream SubExp (StreamForm lore) (Lambda lore) [Input] - | Scatter SubExp (Lambda lore) [Input] [(SubExp, Int, VName)] + = Stream SubExp (StreamForm lore) (Lambda lore) [SubExp] [Input] + | Scatter SubExp (Lambda lore) [Input] [(Shape, Int, VName)] | Screma SubExp (ScremaForm lore) [Input] | Hist SubExp [HistOp lore] (Lambda lore) [Input] deriving (Eq, Show) @@ -404,15 +403,15 @@ -- | Returns the inputs used in a SOAC. inputs :: SOAC lore -> [Input] -inputs (Stream _ _ _ arrs) = arrs +inputs (Stream _ _ _ _ arrs) = arrs inputs (Scatter _len _lam ivs _as) = ivs inputs (Screma _ _ arrs) = arrs inputs (Hist _ _ _ inps) = inps -- | Set the inputs to a SOAC. setInputs :: [Input] -> SOAC lore -> SOAC lore -setInputs arrs (Stream w form lam _) = - Stream (newWidth arrs w) form lam arrs +setInputs arrs (Stream w form lam nes _) = + Stream (newWidth arrs w) form lam nes arrs setInputs arrs (Scatter w lam _ivs as) = Scatter (newWidth arrs w) lam arrs as setInputs arrs (Screma w form _) = @@ -426,15 +425,15 @@ -- | The lambda used in a given SOAC. lambda :: SOAC lore -> Lambda lore -lambda (Stream _ _ lam _) = lam +lambda (Stream _ _ lam _ _) = lam lambda (Scatter _len lam _ivs _as) = lam lambda (Screma _ (ScremaForm _ _ lam) _) = lam lambda (Hist _ _ lam _) = lam -- | Set the lambda used in the SOAC. setLambda :: Lambda lore -> SOAC lore -> SOAC lore -setLambda lam (Stream w form _ arrs) = - Stream w form lam arrs +setLambda lam (Stream w form _ nes arrs) = + Stream w form lam nes arrs setLambda lam (Scatter len _lam ivs as) = Scatter len lam ivs as setLambda lam (Screma w (ScremaForm scan red _) arrs) = @@ -444,20 +443,19 @@ -- | The return type of a SOAC. typeOf :: SOAC lore -> [Type] -typeOf (Stream w form lam _) = - let nes = getStreamAccums form - accrtps = take (length nes) $ lambdaReturnType lam +typeOf (Stream w _ lam nes _) = + let accrtps = take (length nes) $ lambdaReturnType lam arrtps = [ arrayOf (stripArray 1 t) (Shape [w]) NoUniqueness | t <- drop (length nes) (lambdaReturnType lam) ] in accrtps ++ arrtps typeOf (Scatter _w lam _ivs dests) = - zipWith arrayOfRow (drop (n `div` 2) lam_ts) aws + zipWith arrayOfShape val_ts ws where - lam_ts = lambdaReturnType lam - n = length lam_ts - (aws, _, _) = unzip3 dests + indexes = sum $ zipWith (*) ns $ map length ws + val_ts = drop indexes $ lambdaReturnType lam + (ws, ns, _) = unzip3 dests typeOf (Screma w form _) = scremaType w form typeOf (Hist _ ops _ _) = do @@ -467,7 +465,7 @@ -- | The "width" of a SOAC is the expected outer size of its array -- inputs _after_ input-transforms have been carried out. width :: SOAC lore -> SubExp -width (Stream w _ _ _) = w +width (Stream w _ _ _ _) = w width (Scatter len _lam _ivs _as) = len width (Screma w _ _) = w width (Hist w _ _ _) = w @@ -484,8 +482,8 @@ MonadBinder m => SOAC (Lore m) -> m (Futhark.SOAC (Lore m)) -toSOAC (Stream w form lam inps) = - Futhark.Stream w form lam <$> inputsToSubExps inps +toSOAC (Stream w form lam nes inps) = + Futhark.Stream w form lam nes <$> inputsToSubExps inps toSOAC (Scatter len lam ivs dests) = do ivs' <- inputsToSubExps ivs return $ Futhark.Scatter len lam ivs' dests @@ -508,8 +506,8 @@ (Op lore ~ Futhark.SOAC lore, HasScope lore m) => Exp lore -> m (Either NotSOAC (SOAC lore)) -fromExp (Op (Futhark.Stream w form lam as)) = - Right . Stream w form lam <$> traverse varInput as +fromExp (Op (Futhark.Stream w form lam nes as)) = + Right . Stream w form lam nes <$> traverse varInput as fromExp (Op (Futhark.Scatter len lam ivs as)) = Right <$> (Scatter len lam <$> traverse varInput ivs <*> pure as) fromExp (Op (Futhark.Screma w form arrs)) = @@ -553,7 +551,7 @@ strmlam = Lambda strmpar strmbdy loutps empty_lam = Lambda [] (mkBody mempty []) [] -- map(f,a) creates a stream with NO accumulators - return (Stream w (Parallel Disorder Commutative empty_lam []) strmlam inps, []) + return (Stream w (Parallel Disorder Commutative empty_lam) strmlam [] inps, []) | Just (scans, _) <- Futhark.isScanomapSOAC form, Futhark.Scan scan_lam nes <- Futhark.singleScan scans -> do -- scanomap(scan_lam,nes,map_lam,a) => is translated in strem's body to: @@ -643,7 +641,7 @@ strmpar = chunk_param : inpacc_ids ++ strm_inpids strmlam = Lambda strmpar strmbdy (accrtps ++ loutps) return - ( Stream w (Sequential nes) strmlam inps, + ( Stream w Sequential strmlam nes inps, map paramIdent inpacc_ids ) | Just (reds, _) <- Futhark.isRedomapSOAC form, @@ -682,7 +680,7 @@ strmpar = chunk_param : inpacc_ids ++ strm_inpids strmlam = Lambda strmpar strmbdy (accrtps ++ loutps') lam0 <- renameLambda lamin - return (Stream w (Parallel InOrder comm lam0 nes) strmlam inps, []) + return (Stream w (Parallel InOrder comm lam0) strmlam nes inps, []) -- Otherwise it cannot become a stream. _ -> return (soac, []) 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 @@ -1,4 +1,5 @@ {-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE OverloadedStrings #-} {-# OPTIONS_GHC -fno-warn-redundant-constraints #-} @@ -240,27 +241,35 @@ -- | Construct a typed expression from an integer. fromInteger' :: Integer -> TPrimExp t v + -- | Construct a numeric expression from a boolean expression. This + -- can be used to encode arithmetic control flow. + fromBoolExp :: TPrimExp Bool v -> TPrimExp t v + -- | The class of integer types that can be used for constructing -- 'TPrimExp's. class NumExp t => IntExp t instance NumExp Int8 where fromInteger' = isInt8 . ValueExp . IntValue . Int8Value . fromInteger + fromBoolExp = isInt8 . ConvOpExp (BToI Int8) . untyped instance IntExp Int8 instance NumExp Int16 where fromInteger' = isInt16 . ValueExp . IntValue . Int16Value . fromInteger + fromBoolExp = isInt16 . ConvOpExp (BToI Int16) . untyped instance IntExp Int16 instance NumExp Int32 where fromInteger' = isInt32 . ValueExp . IntValue . Int32Value . fromInteger + fromBoolExp = isInt32 . ConvOpExp (BToI Int32) . untyped instance IntExp Int32 instance NumExp Int64 where fromInteger' = isInt64 . ValueExp . IntValue . Int64Value . fromInteger + fromBoolExp = isInt64 . ConvOpExp (BToI Int64) . untyped instance IntExp Int64 @@ -271,10 +280,12 @@ fromRational' :: Rational -> TPrimExp t v instance NumExp Float where - fromInteger' = TPrimExp . ValueExp . FloatValue . Float32Value . fromInteger + fromInteger' = isF32 . ValueExp . FloatValue . Float32Value . fromInteger + fromBoolExp = isF32 . ConvOpExp (SIToFP Int32 Float32) . ConvOpExp (BToI Int32) . untyped instance NumExp Double where fromInteger' = TPrimExp . ValueExp . FloatValue . Float64Value . fromInteger + fromBoolExp = isF64 . ConvOpExp (SIToFP Int32 Float64) . ConvOpExp (BToI Int32) . untyped instance FloatExp Float where fromRational' = TPrimExp . ValueExp . FloatValue . Float32Value . fromRational @@ -327,6 +338,42 @@ | otherwise = numBad "/" (x, y) fromRational = fromRational' + +instance Pretty v => Floating (TPrimExp Float v) where + x ** y = isF32 $ BinOpExp (FPow Float32) (untyped x) (untyped y) + pi = isF32 $ ValueExp $ FloatValue $ Float32Value pi + exp x = isF32 $ FunExp "exp32" [untyped x] $ FloatType Float32 + log x = isF32 $ FunExp "log32" [untyped x] $ FloatType Float32 + sin x = isF32 $ FunExp "sin32" [untyped x] $ FloatType Float32 + cos x = isF32 $ FunExp "cos32" [untyped x] $ FloatType Float32 + tan x = isF32 $ FunExp "tan32" [untyped x] $ FloatType Float32 + asin x = isF32 $ FunExp "asin32" [untyped x] $ FloatType Float32 + acos x = isF32 $ FunExp "acos32" [untyped x] $ FloatType Float32 + atan x = isF32 $ FunExp "atan32" [untyped x] $ FloatType Float32 + sinh x = isF32 $ FunExp "sinh32" [untyped x] $ FloatType Float32 + cosh x = isF32 $ FunExp "cosh32" [untyped x] $ FloatType Float32 + tanh x = isF32 $ FunExp "tanh32" [untyped x] $ FloatType Float32 + asinh x = isF32 $ FunExp "asinh32" [untyped x] $ FloatType Float32 + acosh x = isF32 $ FunExp "acosh32" [untyped x] $ FloatType Float32 + atanh x = isF32 $ FunExp "atanh32" [untyped x] $ FloatType Float32 + +instance Pretty v => Floating (TPrimExp Double v) where + x ** y = isF64 $ BinOpExp (FPow Float64) (untyped x) (untyped y) + pi = isF64 $ ValueExp $ FloatValue $ Float64Value pi + exp x = isF64 $ FunExp "exp64" [untyped x] $ FloatType Float64 + log x = isF64 $ FunExp "log64" [untyped x] $ FloatType Float64 + sin x = isF64 $ FunExp "sin64" [untyped x] $ FloatType Float64 + cos x = isF64 $ FunExp "cos64" [untyped x] $ FloatType Float64 + tan x = isF64 $ FunExp "tan64" [untyped x] $ FloatType Float64 + asin x = isF64 $ FunExp "asin64" [untyped x] $ FloatType Float64 + acos x = isF64 $ FunExp "acos64" [untyped x] $ FloatType Float64 + atan x = isF64 $ FunExp "atan64" [untyped x] $ FloatType Float64 + sinh x = isF64 $ FunExp "sinh64" [untyped x] $ FloatType Float64 + cosh x = isF64 $ FunExp "cosh64" [untyped x] $ FloatType Float64 + tanh x = isF64 $ FunExp "tanh64" [untyped x] $ FloatType Float64 + asinh x = isF64 $ FunExp "asinh64" [untyped x] $ FloatType Float64 + acosh x = isF64 $ FunExp "acosh64" [untyped x] $ FloatType Float64 + atanh x = isF64 $ FunExp "atanh64" [untyped x] $ FloatType Float64 instance (IntExp t, Pretty v) => IntegralExp (TPrimExp t v) where TPrimExp x `div` TPrimExp y 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 @@ -5,8 +5,9 @@ -- | @futhark bench@ module Futhark.CLI.Bench (main) where +import Control.Exception import Control.Monad -import Control.Monad.Except +import Control.Monad.Except hiding (throwError) import qualified Data.ByteString.Char8 as SBS import qualified Data.ByteString.Lazy.Char8 as LBS import Data.Either @@ -176,7 +177,12 @@ | null runner = (binpath, extra_options) | otherwise = (runner, binpath : extra_options) - liftIO $ withServer to_run to_run_args f + liftIO $ withServer to_run to_run_args f `catch` onError + where + onError :: SomeException -> IO a + onError e = do + hPrint stderr e + exitFailure runBenchmark :: BenchOptions -> FutharkExe -> (FilePath, [InputOutputs]) -> IO [BenchResult] runBenchmark opts futhark (program, cases) = do 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 @@ -374,6 +374,11 @@ (NoArg $ Right $ changeFutharkConfig $ \opts -> opts {futharkWerror = True}) "Treat warnings as errors.", Option + "w" + [] + (NoArg $ Right $ changeFutharkConfig $ \opts -> opts {futharkWarn = False}) + "Disable all warnings.", + Option "t" ["type-check"] ( NoArg $ @@ -383,6 +388,13 @@ "Print on standard output the type-checked program.", Option [] + ["no-check"] + ( NoArg $ + Right $ changeFutharkConfig $ \opts -> opts {futharkTypeCheck = False} + ) + "Disable type-checking.", + Option + [] ["pretty-print"] ( NoArg $ Right $ \opts -> @@ -692,7 +704,7 @@ pipeline_config = PipelineConfig { pipelineVerbose = fst (futharkVerbose $ futharkConfig config) > NotVerbose, - pipelineValidate = True + pipelineValidate = futharkTypeCheck $ futharkConfig config } runPolyPass :: 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 @@ -4,10 +4,12 @@ -- | @futhark literate@ module Futhark.CLI.Literate (main) where +import qualified Codec.BMP as BMP import Control.Monad.Except import Data.Bifunctor (bimap, first, second) import Data.Bits -import qualified Data.ByteString.Char8 as BS +import qualified Data.ByteString as BS +import qualified Data.ByteString.Lazy as LBS import Data.Char import Data.Functor import Data.Int (Int64) @@ -19,7 +21,9 @@ import qualified Data.Text.Encoding as T import qualified Data.Text.IO as T import qualified Data.Vector.Storable as SVec +import qualified Data.Vector.Storable.ByteString as SVec import Data.Void +import Data.Word (Word32, Word8) import Futhark.Script import Futhark.Server import Futhark.Test @@ -42,21 +46,21 @@ import Text.Megaparsec.Char import Text.Printf -data AnimParams = AnimParams - { animFPS :: Maybe Int, - animLoop :: Maybe Bool, - animAutoplay :: Maybe Bool, - animFormat :: Maybe T.Text +data VideoParams = VideoParams + { videoFPS :: Maybe Int, + videoLoop :: Maybe Bool, + videoAutoplay :: Maybe Bool, + videoFormat :: Maybe T.Text } deriving (Show) -defaultAnimParams :: AnimParams -defaultAnimParams = - AnimParams - { animFPS = Nothing, - animLoop = Nothing, - animAutoplay = Nothing, - animFormat = Nothing +defaultVideoParams :: VideoParams +defaultVideoParams = + VideoParams + { videoFPS = Nothing, + videoLoop = Nothing, + videoAutoplay = Nothing, + videoFormat = Nothing } data Directive @@ -66,7 +70,7 @@ | DirectiveImg Exp | DirectivePlot Exp (Maybe (Int, Int)) | DirectiveGnuplot Exp T.Text - | DirectiveAnim Exp AnimParams + | DirectiveVideo Exp VideoParams deriving (Show) varsInDirective :: Directive -> S.Set EntryName @@ -76,7 +80,7 @@ varsInDirective (DirectiveImg e) = varsInExp e varsInDirective (DirectivePlot e _) = varsInExp e varsInDirective (DirectiveGnuplot e _) = varsInExp e -varsInDirective (DirectiveAnim e _) = varsInExp e +varsInDirective (DirectiveVideo e _) = varsInExp e pprDirective :: Bool -> Directive -> PP.Doc pprDirective _ (DirectiveRes e) = @@ -100,18 +104,18 @@ map PP.strictText (T.lines script) pprDirective False (DirectiveGnuplot e _) = "> :gnuplot " <> PP.align (PP.ppr e) -pprDirective False (DirectiveAnim e _) = - "> :anim " <> PP.align (PP.ppr e) -pprDirective True (DirectiveAnim e params) = - "> :anim " <> PP.ppr e +pprDirective False (DirectiveVideo e _) = + "> :video " <> PP.align (PP.ppr e) +pprDirective True (DirectiveVideo e params) = + "> :video " <> PP.ppr e <> if null params' then mempty else PP.stack $ ";" : params' where params' = catMaybes - [ p "fps" animFPS PP.ppr, - p "loop" animLoop ppBool, - p "autoplay" animAutoplay ppBool, - p "format" animFormat PP.strictText + [ p "fps" videoFPS PP.ppr, + p "loop" videoLoop ppBool, + p "autoplay" videoAutoplay ppBool, + p "format" videoFormat PP.strictText ] ppBool b = if b then "true" else "false" p s f ppr = do @@ -176,10 +180,10 @@ *> token "(" *> ((,) <$> parseInt <* token "," <*> parseInt) <* token ")" -parseAnimParams :: Parser AnimParams -parseAnimParams = - fmap (fromMaybe defaultAnimParams) $ - optional $ ";" *> hspace *> eol *> "-- " *> parseParams defaultAnimParams +parseVideoParams :: Parser VideoParams +parseVideoParams = + fmap (fromMaybe defaultVideoParams) $ + optional $ ";" *> hspace *> eol *> "-- " *> parseParams defaultVideoParams where parseParams params = choice @@ -192,24 +196,24 @@ pLoop params = do token "loop:" b <- parseBool - pure params {animLoop = Just b} + pure params {videoLoop = Just b} pFPS params = do token "fps:" fps <- parseInt - pure params {animFPS = Just fps} + pure params {videoFPS = Just fps} pAutoplay params = do token "autoplay:" b <- parseBool - pure params {animAutoplay = Just b} + pure params {videoAutoplay = Just b} pFormat params = do token "format:" s <- lexeme $ takeWhileP Nothing (not . isSpace) - pure params {animFormat = Just s} + pure params {videoFormat = Just s} parseBlock :: Parser Block parseBlock = choice - [ token "-- >" $> BlockDirective <*> parseDirective <* void eol, + [ token "-- >" $> BlockDirective <*> parseDirective, BlockCode <$> parseTestBlock, BlockCode <$> parseBlockCode, BlockComment <$> parseBlockComment @@ -223,16 +227,16 @@ directiveName "brief" $> DirectiveBrief <*> parseDirective, directiveName "img" $> DirectiveImg - <*> parseExp postlexeme, + <*> parseExp postlexeme <* eol, directiveName "plot2d" $> DirectivePlot <*> parseExp postlexeme - <*> parsePlotParams, + <*> parsePlotParams <* eol, directiveName "gnuplot" $> DirectiveGnuplot <*> parseExp postlexeme <*> (";" *> hspace *> eol *> parseBlockComment), - directiveName "anim" $> DirectiveAnim + (directiveName "video" <|> directiveName "video") $> DirectiveVideo <*> parseExp postlexeme - <*> parseAnimParams + <*> parseVideoParams <* eol ] directiveName s = try $ token (":" <> s) @@ -264,57 +268,46 @@ join . liftIO . withSystemTempDirectory "futhark-literate" $ \dir -> either throwError pure <$> runExceptT (f dir) -ppmHeader :: Int -> Int -> BS.ByteString -ppmHeader h w = - "P6\n" <> BS.pack (show w) <> " " <> BS.pack (show h) <> "\n255\n" - -rgbIntToImg :: - (Integral a, Bits a, SVec.Storable a) => - Int -> - Int -> - SVec.Vector a -> - BS.ByteString -rgbIntToImg h w bytes = - ppmHeader h w <> fst (BS.unfoldrN (h * w * 3) byte 0) - where - getChan word chan = - (word `shiftR` (chan * 8)) .&. 0xFF - byte i = - Just - ( chr . max 0 . fromIntegral $ - getChan (bytes SVec.! (i `div` 3)) (2 - (i `mod` 3)), - i + 1 - ) - greyFloatToImg :: (RealFrac a, SVec.Storable a) => - Int -> - Int -> SVec.Vector a -> - BS.ByteString -greyFloatToImg h w bytes = - ppmHeader h w <> fst (BS.unfoldrN (h * w * 3) byte 0) + SVec.Vector Word32 +greyFloatToImg = SVec.map grey where - byte i = - Just (chr . max 0 $ round (bytes SVec.! (i `div` 3)) * 255, i + 1) + grey i = + let i' = round (i * 255) .&. 0xFF + in (i' `shiftL` 16) .|. (i' `shiftL` 8) .|. i' -valueToPPM :: Value -> Maybe BS.ByteString -valueToPPM v@(Word32Value _ bytes) +-- BMPs are RGBA and bottom-up where we assumes images are top-down +-- and ARGB. We fix this up before encoding the BMP. This is +-- probably a little slower than it has to be. +vecToBMP :: Int -> Int -> SVec.Vector Word32 -> LBS.ByteString +vecToBMP h w = BMP.renderBMP . BMP.packRGBA32ToBMP24 w h . SVec.vectorToByteString . frobVec + where + frobVec vec = SVec.generate (h * w * 4) (pix vec) + pix vec l = + let (i, j) = (l `div` 4) `divMod` w + argb = vec SVec.! ((h -1 - i) * w + j) + c = (argb `shiftR` (24 - ((l + 1) `mod` 4) * 8)) .&. 0xFF + in fromIntegral c :: Word8 + +valueToBMP :: Value -> Maybe LBS.ByteString +valueToBMP v@(Word32Value _ bytes) | [h, w] <- valueShape v = - Just $ rgbIntToImg h w bytes -valueToPPM v@(Int32Value _ bytes) + Just $ vecToBMP h w bytes +valueToBMP v@(Int32Value _ bytes) | [h, w] <- valueShape v = - Just $ rgbIntToImg h w bytes -valueToPPM v@(Float32Value _ bytes) + Just $ vecToBMP h w $ SVec.map fromIntegral bytes +valueToBMP v@(Float32Value _ bytes) | [h, w] <- valueShape v = - Just $ greyFloatToImg h w bytes -valueToPPM v@(Float64Value _ bytes) + Just $ vecToBMP h w $ greyFloatToImg bytes +valueToBMP v@(Float64Value _ bytes) | [h, w] <- valueShape v = - Just $ greyFloatToImg h w bytes -valueToPPM _ = Nothing + Just $ vecToBMP h w $ greyFloatToImg bytes +valueToBMP _ = Nothing -valueToPPMs :: Value -> Maybe [BS.ByteString] -valueToPPMs = mapM valueToPPM . valueElems +valueToBMPs :: Value -> Maybe [LBS.ByteString] +valueToBMPs = mapM valueToBMP . valueElems system :: FilePath -> [String] -> T.Text -> ScriptM T.Text system prog options input = do @@ -333,12 +326,12 @@ where prog' = "'" <> T.pack prog <> "'" -ppmToPNG :: FilePath -> ScriptM FilePath -ppmToPNG ppm = do - void $ system "convert" [ppm, png] mempty +bmpToPNG :: FilePath -> ScriptM FilePath +bmpToPNG bmp = do + void $ system "convert" [bmp, png] mempty pure png where - png = ppm `replaceExtension` "png" + png = bmp `replaceExtension` "png" formatDataForGnuplot :: [Value] -> T.Text formatDataForGnuplot = T.unlines . map line . transpose . map valueElems @@ -348,7 +341,7 @@ imgBlock :: FilePath -> T.Text imgBlock f = "\n\n![](" <> T.pack f <> ")\n\n" -videoBlock :: AnimParams -> FilePath -> T.Text +videoBlock :: VideoParams -> FilePath -> T.Text videoBlock opts f = "\n\n![](" <> T.pack f <> ")" <> opts' <> "\n\n" where opts' @@ -361,8 +354,8 @@ if b then s <> "=\"true\"" else s <> "=\"false\"" | otherwise = mempty - loop = boolOpt "loop" animLoop - autoplay = boolOpt "autoplay" animAutoplay + loop = boolOpt "loop" videoLoop + autoplay = boolOpt "autoplay" videoAutoplay plottable :: CompoundValue -> Maybe [Value] plottable (ValueTuple vs) = do @@ -386,13 +379,54 @@ liftIO $ T.writeFile fname $ formatDataForGnuplot vs withGnuplotData ((f, f <> "='" <> T.pack fname <> "'") : sets) xys cont +loadBMP :: FilePath -> ScriptM (Compound Value) +loadBMP bmpfile = do + res <- liftIO $ BMP.readBMP bmpfile + case res of + Left err -> + throwError $ "Failed to read BMP:\n" <> T.pack (show err) + Right bmp -> do + let bmp_bs = BMP.unpackBMPToRGBA32 bmp + (w, h) = BMP.bmpDimensions bmp + shape = SVec.fromList [fromIntegral h, fromIntegral w] + pix l = + let (i, j) = l `divMod` w + l' = (h -1 - i) * w + j + r = fromIntegral $ bmp_bs `BS.index` (l' * 4) + g = fromIntegral $ bmp_bs `BS.index` (l' * 4 + 1) + b = fromIntegral $ bmp_bs `BS.index` (l' * 4 + 2) + a = fromIntegral $ bmp_bs `BS.index` (l' * 4 + 3) + in (a `shiftL` 24) .|. (r `shiftL` 16) .|. (g `shiftL` 8) .|. b + pure $ ValueAtom $ Word32Value shape $ SVec.generate (w * h) pix + +loadImage :: FilePath -> ScriptM (Compound Value) +loadImage imgfile = + withTempDir $ \dir -> do + let bmpfile = dir imgfile `replaceExtension` "bmp" + void $ system "convert" [imgfile, "-type", "TrueColorAlpha", bmpfile] mempty + loadBMP bmpfile + +literateBuiltin :: EvalBuiltin ScriptM +literateBuiltin "loadimg" vs = + case vs of + [ValueAtom v] + | Just path <- getValue v -> do + let path' = map (chr . fromIntegral) (path :: [Word8]) + loadImage path' + _ -> + throwError $ + "$imgfile does not accept arguments of types: " + <> T.intercalate ", " (map (prettyText . fmap valueType) vs) +literateBuiltin f _ = + throwError $ "Unknown builtin function $" <> prettyText f + processDirective :: FilePath -> ScriptServer -> Int -> Directive -> ScriptM T.Text processDirective imgdir server i (DirectiveBrief d) = processDirective imgdir server i d processDirective imgdir server i (DirectiveCovert d) = processDirective imgdir server i d processDirective _ server _ (DirectiveRes e) = do - vs <- evalExpToGround server e + vs <- evalExpToGround literateBuiltin server e pure $ T.unlines [ "", @@ -403,15 +437,15 @@ ] -- processDirective imgdir server i (DirectiveImg e) = do - vs <- evalExpToGround server e + vs <- evalExpToGround literateBuiltin server e case vs of ValueAtom v - | Just ppm <- valueToPPM v -> do - let ppmfile = imgdir "img" <> show i <.> ".ppm" + | Just bmp <- valueToBMP v -> do + let bmpfile = imgdir "img" <> show i <.> ".bmp" liftIO $ createDirectoryIfMissing True imgdir - liftIO $ BS.writeFile ppmfile ppm - pngfile <- ppmToPNG ppmfile - liftIO $ removeFile ppmfile + liftIO $ LBS.writeFile bmpfile bmp + pngfile <- bmpToPNG bmpfile + liftIO $ removeFile bmpfile pure $ imgBlock pngfile _ -> throwError $ @@ -419,7 +453,7 @@ <> prettyText (fmap valueType vs) -- processDirective imgdir server i (DirectivePlot e size) = do - v <- evalExpToGround server e + v <- evalExpToGround literateBuiltin server e case v of _ | Just vs <- plottable2d v -> @@ -464,7 +498,7 @@ pure $ imgBlock pngfile -- processDirective imgdir server i (DirectiveGnuplot e script) = do - vs <- evalExpToGround server e + vs <- evalExpToGround literateBuiltin server e case vs of ValueRecord m | Just m' <- traverse plottable m -> @@ -487,23 +521,23 @@ void $ system "gnuplot" [] script' pure $ imgBlock pngfile -- -processDirective imgdir server i (DirectiveAnim e params) = do +processDirective imgdir server i (DirectiveVideo e params) = do when (format `notElem` ["webm", "gif"]) $ - throwError $ "Unknown animation format: " <> format + throwError $ "Unknown video format: " <> format - v <- evalExp server e + v <- evalExp literateBuiltin server e let nope = throwError $ - "Cannot animate value of type " <> prettyText (fmap scriptValueType v) + "Cannot produce video from value of type " <> prettyText (fmap scriptValueType v) case v of ValueAtom SValue {} -> do ValueAtom arr <- getExpValue server v - case valueToPPMs arr of + case valueToBMPs arr of Nothing -> nope - Just ppms -> + Just bmps -> withTempDir $ \dir -> do - zipWithM_ (writePPMFile dir) [0 ..] ppms - ppmsToVideo dir + zipWithM_ (writeBMPFile dir) [0 ..] bmps + bmpsToVideo dir ValueTuple [stepfun, initial, num_frames] | ValueAtom (SFun stepfun' _ [_, _] closure) <- stepfun, ValueAtom (SValue _ _) <- initial, @@ -513,28 +547,32 @@ withTempDir $ \dir -> do let num_frames_int = fromIntegral (num_frames' :: Int64) renderFrames dir (stepfun', map ValueAtom closure) initial num_frames_int - ppmsToVideo dir + bmpsToVideo dir _ -> nope - when (animFormat params == Just "gif") $ do + when (videoFormat params == Just "gif") $ do void $ system "ffmpeg" ["-i", webmfile, giffile] mempty liftIO $ removeFile webmfile - pure $ videoBlock params animfile + pure $ videoBlock params videofile where - framerate = fromMaybe 30 $ animFPS params - format = fromMaybe "webm" $ animFormat params - webmfile = imgdir "anim" <> show i <.> "webm" - giffile = imgdir "anim" <> show i <.> "gif" - ppmfile dir j = dir printf "frame%010d.ppm" (j :: Int) - animfile = imgdir "anim" <> show i <.> T.unpack format + framerate = fromMaybe 30 $ videoFPS params + format = fromMaybe "webm" $ videoFormat params + webmfile = imgdir "video" <> show i <.> "webm" + giffile = imgdir "video" <> show i <.> "gif" + bmpfile dir j = dir printf "frame%010d.bmp" (j :: Int) + videofile = imgdir "video" <> show i <.> T.unpack format renderFrames dir (stepfun, closure) initial num_frames = foldM_ frame initial [0 .. num_frames -1] where frame old_state j = do - v <- evalExp server . Call stepfun . map valueToExp $ closure ++ [old_state] + v <- + evalExp literateBuiltin server + . Call (FuncFut stepfun) + . map valueToExp + $ closure ++ [old_state] freeValue server old_state let nope = @@ -546,14 +584,14 @@ ValueTuple [arr_v@(ValueAtom SValue {}), new_state] -> do ValueAtom arr <- getExpValue server arr_v freeValue server arr_v - case valueToPPM arr of + case valueToBMP arr of Nothing -> nope - Just ppm -> do - writePPMFile dir j ppm + Just bmp -> do + writeBMPFile dir j bmp pure new_state _ -> nope - ppmsToVideo dir = do + bmpsToVideo dir = do liftIO $ createDirectoryIfMissing True imgdir void $ system @@ -562,7 +600,7 @@ "-r", show framerate, "-i", - dir "frame%010d.ppm", + dir "frame%010d.bmp", "-c:v", "libvpx-vp9", "-pix_fmt", @@ -573,9 +611,8 @@ ] mempty - writePPMFile dir j ppm = do - let fname = ppmfile dir j - liftIO $ BS.writeFile fname ppm + writeBMPFile dir j bmp = + liftIO $ LBS.writeFile (bmpfile dir j) bmp -- Did this script block succeed or fail? data Failure = Failure | Success 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 @@ -259,6 +259,7 @@ cudaSizeClass SizeGroup = "block_size" cudaSizeClass SizeNumGroups = "grid_size" cudaSizeClass SizeTile = "tile_size" + cudaSizeClass SizeRegTile = "reg_tile_size" cudaSizeClass SizeLocalMemory = "shared_memory" cudaSizeClass (SizeBespoke x _) = pretty x callKernel (LaunchKernel safety kernel_name args num_blocks block_size) = do diff --git a/src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs b/src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs --- a/src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs +++ b/src/Futhark/CodeGen/Backends/CCUDA/Boilerplate.hs @@ -91,33 +91,11 @@ let size_name_inits = map (\k -> [C.cinit|$string:(pretty k)|]) $ M.keys sizes size_var_inits = map (\k -> [C.cinit|$string:(zEncodeString (pretty k))|]) $ M.keys sizes size_class_inits = map (\c -> [C.cinit|$string:(pretty c)|]) $ M.elems sizes - num_sizes = M.size sizes GC.earlyDecl [C.cedecl|static const char *size_names[] = { $inits:size_name_inits };|] GC.earlyDecl [C.cedecl|static const char *size_vars[] = { $inits:size_var_inits };|] GC.earlyDecl [C.cedecl|static const char *size_classes[] = { $inits:size_class_inits };|] - GC.publicDef_ "get_num_sizes" GC.InitDecl $ \s -> - ( [C.cedecl|int $id:s(void);|], - [C.cedecl|int $id:s(void) { - return $int:num_sizes; - }|] - ) - - GC.publicDef_ "get_size_name" GC.InitDecl $ \s -> - ( [C.cedecl|const char* $id:s(int);|], - [C.cedecl|const char* $id:s(int i) { - return size_names[i]; - }|] - ) - - GC.publicDef_ "get_size_class" GC.InitDecl $ \s -> - ( [C.cedecl|const char* $id:s(int);|], - [C.cedecl|const char* $id:s(int i) { - return size_classes[i]; - }|] - ) - generateConfigFuns :: M.Map Name SizeClass -> GC.CompilerM OpenCL () String generateConfigFuns sizes = do let size_decls = map (\k -> [C.csdecl|typename int64_t $id:k;|]) $ M.keys sizes @@ -255,6 +233,13 @@ }|] ) + GC.publicDef_ "context_config_set_default_reg_tile_size" GC.InitDecl $ \s -> + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int num);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int size) { + cfg->cu_cfg.default_reg_tile_size = size; + }|] + ) + GC.publicDef_ "context_config_set_default_threshold" GC.InitDecl $ \s -> ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int num);|], [C.cedecl|void $id:s(struct $id:cfg* cfg, int size) { @@ -292,6 +277,12 @@ cfg->cu_cfg.default_tile_size = size_value; return 0; } + + if (strcmp(size_name, "default_reg_tile_size") == 0) { + cfg->cu_cfg.default_reg_tile_size = size_value; + return 0; + } + return 1; }|] ) diff --git a/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs b/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs --- a/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs +++ b/src/Futhark/CodeGen/Backends/COpenCL/Boilerplate.hs @@ -93,27 +93,6 @@ GC.earlyDecl [C.cedecl|static const char *size_vars[] = { $inits:size_var_inits };|] GC.earlyDecl [C.cedecl|static const char *size_classes[] = { $inits:size_class_inits };|] - GC.publicDef_ "get_num_sizes" GC.InitDecl $ \s -> - ( [C.cedecl|int $id:s(void);|], - [C.cedecl|int $id:s(void) { - return $int:num_sizes; - }|] - ) - - GC.publicDef_ "get_size_name" GC.InitDecl $ \s -> - ( [C.cedecl|const char* $id:s(int);|], - [C.cedecl|const char* $id:s(int i) { - return size_names[i]; - }|] - ) - - GC.publicDef_ "get_size_class" GC.InitDecl $ \s -> - ( [C.cedecl|const char* $id:s(int);|], - [C.cedecl|const char* $id:s(int i) { - return size_classes[i]; - }|] - ) - let size_decls = map (\k -> [C.csdecl|typename int64_t $id:k;|]) $ M.keys sizes GC.earlyDecl [C.cedecl|struct sizes { $sdecls:size_decls };|] cfg <- GC.publicDef "context_config" GC.InitDecl $ \s -> @@ -267,6 +246,13 @@ }|] ) + GC.publicDef_ "context_config_set_default_reg_tile_size" GC.InitDecl $ \s -> + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int num);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int size) { + cfg->opencl.default_reg_tile_size = size; + }|] + ) + GC.publicDef_ "context_config_set_default_threshold" GC.InitDecl $ \s -> ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int num);|], [C.cedecl|void $id:s(struct $id:cfg* cfg, int size) { @@ -305,6 +291,11 @@ return 0; } + if (strcmp(size_name, "default_reg_tile_size") == 0) { + cfg->opencl.default_reg_tile_size = size_value; + return 0; + } + return 1; }|] ) @@ -648,6 +639,7 @@ NumGroups -> [C.cexp|ctx->cfg.default_num_groups|] GroupSize -> [C.cexp|ctx->cfg.default_group_size|] TileSize -> [C.cexp|ctx->cfg.default_tile_size|] + RegTileSize -> [C.cexp|ctx->cfg.default_reg_tile_size|] Threshold -> [C.cexp|ctx->cfg.default_threshold|] get_size = @@ -707,58 +699,17 @@ optionAction = [C.cstm|futhark_context_config_set_default_tile_size(cfg, atoi(optarg));|] }, Option - { optionLongName = "default-threshold", + { optionLongName = "default-reg-tile-size", optionShortName = Nothing, optionArgument = RequiredArgument "INT", - optionDescription = "The default parallelism threshold.", - optionAction = [C.cstm|futhark_context_config_set_default_threshold(cfg, atoi(optarg));|] - }, - Option - { optionLongName = "print-sizes", - optionShortName = Nothing, - optionArgument = NoArgument, - optionDescription = "Print all sizes that can be set with -size or --tuning.", - optionAction = - [C.cstm|{ - int n = futhark_get_num_sizes(); - for (int i = 0; i < n; i++) { - printf("%s (%s)\n", futhark_get_size_name(i), - futhark_get_size_class(i)); - } - exit(0); - }|] - }, - Option - { optionLongName = "size", - optionShortName = Nothing, - optionArgument = RequiredArgument "ASSIGNMENT", - optionDescription = "Set a configurable run-time parameter to the given value.", - optionAction = - [C.cstm|{ - char *name = optarg; - char *equals = strstr(optarg, "="); - char *value_str = equals != NULL ? equals+1 : optarg; - int value = atoi(value_str); - if (equals != NULL) { - *equals = 0; - if (futhark_context_config_set_size(cfg, name, value) != 0) { - futhark_panic(1, "Unknown size: %s\n", name); - } - } else { - futhark_panic(1, "Invalid argument for size option: %s\n", optarg); - }}|] + optionDescription = "The default register tile size used when performing two-dimensional tiling.", + optionAction = [C.cstm|futhark_context_config_set_default_reg_tile_size(cfg, atoi(optarg));|] }, Option - { optionLongName = "tuning", + { optionLongName = "default-threshold", optionShortName = Nothing, - optionArgument = RequiredArgument "FILE", - optionDescription = "Read size=value assignments from the given file.", - optionAction = - [C.cstm|{ - char *ret = load_tuning_file(optarg, cfg, (int(*)(void*, const char*, size_t)) - futhark_context_config_set_size); - if (ret != NULL) { - futhark_panic(1, "When loading tuning from '%s': %s\n", optarg, ret); - }}|] + optionArgument = RequiredArgument "INT", + optionDescription = "The default parallelism threshold.", + optionAction = [C.cstm|futhark_context_config_set_default_threshold(cfg, atoi(optarg));|] } ] diff --git a/src/Futhark/CodeGen/Backends/GenericC.hs b/src/Futhark/CodeGen/Backends/GenericC.hs --- a/src/Futhark/CodeGen/Backends/GenericC.hs +++ b/src/Futhark/CodeGen/Backends/GenericC.hs @@ -1463,6 +1463,7 @@ $esc:("#include ") $esc:("#include ") $esc:("#include ") +$esc:("#include ") $esc:header_extra @@ -1540,6 +1541,27 @@ ops <- asks envOperations profilereport <- gets $ DL.toList . compProfileItems + publicDef_ "get_num_sizes" InitDecl $ \s -> + ( [C.cedecl|int $id:s(void);|], + [C.cedecl|int $id:s(void) { + return sizeof(size_names)/sizeof(size_names[0]); + }|] + ) + + publicDef_ "get_size_name" InitDecl $ \s -> + ( [C.cedecl|const char* $id:s(int);|], + [C.cedecl|const char* $id:s(int i) { + return size_names[i]; + }|] + ) + + publicDef_ "get_size_class" InitDecl $ \s -> + ( [C.cedecl|const char* $id:s(int);|], + [C.cedecl|const char* $id:s(int i) { + return size_classes[i]; + }|] + ) + publicDef_ "context_report" MiscDecl $ \s -> ( [C.cedecl|char* $id:s($ty:ctx *ctx);|], [C.cedecl|char* $id:s($ty:ctx *ctx) { @@ -1851,6 +1873,12 @@ compilePrimExp f (UnOpExp USignum {} x) = do x' <- compilePrimExp f x return [C.cexp|($exp:x' > 0) - ($exp:x' < 0) != 0|] +compilePrimExp f (UnOpExp (FSignum Float32) x) = do + x' <- compilePrimExp f x + return [C.cexp|fsignum32($exp:x')|] +compilePrimExp f (UnOpExp (FSignum Float64) x) = do + x' <- compilePrimExp f x + return [C.cexp|fsignum32($exp:x')|] compilePrimExp f (CmpOpExp cmp x y) = do x' <- compilePrimExp f x y' <- compilePrimExp f y diff --git a/src/Futhark/CodeGen/Backends/GenericC/CLI.hs b/src/Futhark/CodeGen/Backends/GenericC/CLI.hs --- a/src/Futhark/CodeGen/Backends/GenericC/CLI.hs +++ b/src/Futhark/CodeGen/Backends/GenericC/CLI.hs @@ -75,6 +75,54 @@ fut_progname, option_descriptions); exit(0); }|] + }, + Option + { optionLongName = "print-sizes", + optionShortName = Nothing, + optionArgument = NoArgument, + optionDescription = "Print all sizes that can be set with --size or --tuning.", + optionAction = + [C.cstm|{ + int n = futhark_get_num_sizes(); + for (int i = 0; i < n; i++) { + printf("%s (%s)\n", futhark_get_size_name(i), + futhark_get_size_class(i)); + } + exit(0); + }|] + }, + Option + { optionLongName = "size", + optionShortName = Nothing, + optionArgument = RequiredArgument "ASSIGNMENT", + optionDescription = "Set a configurable run-time parameter to the given value.", + optionAction = + [C.cstm|{ + char *name = optarg; + char *equals = strstr(optarg, "="); + char *value_str = equals != NULL ? equals+1 : optarg; + int value = atoi(value_str); + if (equals != NULL) { + *equals = 0; + if (futhark_context_config_set_size(cfg, name, value) != 0) { + futhark_panic(1, "Unknown size: %s\n", name); + } + } else { + futhark_panic(1, "Invalid argument for size option: %s\n", optarg); + }}|] + }, + Option + { optionLongName = "tuning", + optionShortName = Nothing, + optionArgument = RequiredArgument "FILE", + optionDescription = "Read size=value assignments from the given file.", + optionAction = + [C.cstm|{ + char *ret = load_tuning_file(optarg, cfg, (int(*)(void*, const char*, size_t)) + futhark_context_config_set_size); + if (ret != NULL) { + futhark_panic(1, "When loading tuning from '%s': %s\n", optarg, ret); + }}|] } ] where diff --git a/src/Futhark/CodeGen/Backends/GenericC/Server.hs b/src/Futhark/CodeGen/Backends/GenericC/Server.hs --- a/src/Futhark/CodeGen/Backends/GenericC/Server.hs +++ b/src/Futhark/CodeGen/Backends/GenericC/Server.hs @@ -48,6 +48,54 @@ fut_progname, option_descriptions); exit(0); }|] + }, + Option + { optionLongName = "print-sizes", + optionShortName = Nothing, + optionArgument = NoArgument, + optionDescription = "Print all sizes that can be set with --size or --tuning.", + optionAction = + [C.cstm|{ + int n = futhark_get_num_sizes(); + for (int i = 0; i < n; i++) { + printf("%s (%s)\n", futhark_get_size_name(i), + futhark_get_size_class(i)); + } + exit(0); + }|] + }, + Option + { optionLongName = "size", + optionShortName = Nothing, + optionArgument = RequiredArgument "ASSIGNMENT", + optionDescription = "Set a configurable run-time parameter to the given value.", + optionAction = + [C.cstm|{ + char *name = optarg; + char *equals = strstr(optarg, "="); + char *value_str = equals != NULL ? equals+1 : optarg; + int value = atoi(value_str); + if (equals != NULL) { + *equals = 0; + if (futhark_context_config_set_size(cfg, name, value) != 0) { + futhark_panic(1, "Unknown size: %s\n", name); + } + } else { + futhark_panic(1, "Invalid argument for size option: %s\n", optarg); + }}|] + }, + Option + { optionLongName = "tuning", + optionShortName = Nothing, + optionArgument = RequiredArgument "FILE", + optionDescription = "Read size=value assignments from the given file.", + optionAction = + [C.cstm|{ + char *ret = load_tuning_file(optarg, cfg, (int(*)(void*, const char*, size_t)) + futhark_context_config_set_size); + if (ret != NULL) { + futhark_panic(1, "When loading tuning from '%s': %s\n", optarg, ret); + }}|] } ] 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 @@ -989,6 +989,7 @@ FAbs {} -> "abs" SSignum {} -> "ssignum" USignum {} -> "usignum" + FSignum {} -> "np.sign" compileBinOpLike :: Monad m => diff --git a/src/Futhark/CodeGen/Backends/MulticoreC.hs b/src/Futhark/CodeGen/Backends/MulticoreC.hs --- a/src/Futhark/CodeGen/Backends/MulticoreC.hs +++ b/src/Futhark/CodeGen/Backends/MulticoreC.hs @@ -191,6 +191,18 @@ }|] ) + GC.earlyDecl [C.cedecl|static const char *size_names[0];|] + GC.earlyDecl [C.cedecl|static const char *size_vars[0];|] + GC.earlyDecl [C.cedecl|static const char *size_classes[0];|] + + GC.publicDef_ "context_config_set_size" GC.InitDecl $ \s -> + ( [C.cedecl|int $id:s(struct $id:cfg* cfg, const char *size_name, size_t size_value);|], + [C.cedecl|int $id:s(struct $id:cfg* cfg, const char *size_name, size_t size_value) { + (void)cfg; (void)size_name; (void)size_value; + return 1; + }|] + ) + cliOptions :: [Option] cliOptions = [ Option diff --git a/src/Futhark/CodeGen/Backends/PyOpenCL.hs b/src/Futhark/CodeGen/Backends/PyOpenCL.hs --- a/src/Futhark/CodeGen/Backends/PyOpenCL.hs +++ b/src/Futhark/CodeGen/Backends/PyOpenCL.hs @@ -58,6 +58,7 @@ Assign (Var "default_group_size") None, Assign (Var "default_num_groups") None, Assign (Var "default_tile_size") None, + Assign (Var "default_reg_tile_size") None, Assign (Var "fut_opencl_src") $ RawStringLiteral $ opencl_prelude ++ opencl_code ] @@ -80,6 +81,7 @@ "default_group_size=default_group_size", "default_num_groups=default_num_groups", "default_tile_size=default_tile_size", + "default_reg_tile_size=default_reg_tile_size", "default_threshold=default_threshold", "sizes=sizes" ] @@ -126,6 +128,13 @@ optionArgument = RequiredArgument "int", optionAction = [Assign (Var "default_tile_size") $ Var "optarg"] + }, + Option + { optionLongName = "default-reg-tile-size", + optionShortName = Nothing, + optionArgument = RequiredArgument "int", + optionAction = + [Assign (Var "default_reg_tile_size") $ Var "optarg"] }, Option { optionLongName = "size", diff --git a/src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs b/src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs --- a/src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs +++ b/src/Futhark/CodeGen/Backends/PyOpenCL/Boilerplate.hs @@ -55,6 +55,7 @@ default_group_size=default_group_size, default_num_groups=default_num_groups, default_tile_size=default_tile_size, + default_reg_tile_size=default_reg_tile_size, default_threshold=default_threshold, size_heuristics=size_heuristics, required_types=$types', @@ -117,6 +118,7 @@ NumGroups -> String "num_groups" GroupSize -> String "group_size" TileSize -> String "tile_size" + RegTileSize -> String "reg_tile_size" Threshold -> String "threshold" what' = diff --git a/src/Futhark/CodeGen/Backends/SequentialC.hs b/src/Futhark/CodeGen/Backends/SequentialC.hs --- a/src/Futhark/CodeGen/Backends/SequentialC.hs +++ b/src/Futhark/CodeGen/Backends/SequentialC.hs @@ -1,5 +1,3 @@ -{-# LANGUAGE QuasiQuotes #-} - -- | C code generator. This module can convert a correct ImpCode -- program to an equivalent C program. The C code is strictly -- sequential, but can handle the full Futhark language. @@ -14,17 +12,17 @@ import Control.Monad import qualified Futhark.CodeGen.Backends.GenericC as GC +import Futhark.CodeGen.Backends.SequentialC.Boilerplate import qualified Futhark.CodeGen.ImpCode.Sequential as Imp import qualified Futhark.CodeGen.ImpGen.Sequential as ImpGen import Futhark.IR.SeqMem import Futhark.MonadFreshNames -import qualified Language.C.Quote.OpenCL as C -- | Compile the program to sequential C. compileProg :: MonadFreshNames m => Prog SeqMem -> m (ImpGen.Warnings, GC.CParts) compileProg = traverse - (GC.compileProg "c" operations generateContext "" [DefaultSpace] []) + (GC.compileProg "c" operations generateBoilerplate "" [DefaultSpace] []) <=< ImpGen.compileProg where operations :: GC.Operations Imp.Sequential () @@ -32,97 +30,3 @@ GC.defaultOperations { GC.opsCompiler = const $ return () } - - generateContext = do - cfg <- GC.publicDef "context_config" GC.InitDecl $ \s -> - ( [C.cedecl|struct $id:s;|], - [C.cedecl|struct $id:s { int debugging; };|] - ) - - GC.publicDef_ "context_config_new" GC.InitDecl $ \s -> - ( [C.cedecl|struct $id:cfg* $id:s(void);|], - [C.cedecl|struct $id:cfg* $id:s(void) { - struct $id:cfg *cfg = (struct $id:cfg*) malloc(sizeof(struct $id:cfg)); - if (cfg == NULL) { - return NULL; - } - cfg->debugging = 0; - return cfg; - }|] - ) - - GC.publicDef_ "context_config_free" GC.InitDecl $ \s -> - ( [C.cedecl|void $id:s(struct $id:cfg* cfg);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg) { - free(cfg); - }|] - ) - - GC.publicDef_ "context_config_set_debugging" GC.InitDecl $ \s -> - ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg, int detail) { - cfg->debugging = detail; - }|] - ) - - GC.publicDef_ "context_config_set_logging" GC.InitDecl $ \s -> - ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag);|], - [C.cedecl|void $id:s(struct $id:cfg* cfg, int detail) { - /* Does nothing for this backend. */ - (void)cfg; (void)detail; - }|] - ) - - (fields, init_fields) <- GC.contextContents - - ctx <- GC.publicDef "context" GC.InitDecl $ \s -> - ( [C.cedecl|struct $id:s;|], - [C.cedecl|struct $id:s { - int detail_memory; - int debugging; - int profiling; - int logging; - typename lock_t lock; - char *error; - typename FILE *log; - int profiling_paused; - $sdecls:fields - };|] - ) - - GC.publicDef_ "context_new" GC.InitDecl $ \s -> - ( [C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg);|], - [C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg) { - struct $id:ctx* ctx = (struct $id:ctx*) malloc(sizeof(struct $id:ctx)); - if (ctx == NULL) { - return NULL; - } - ctx->detail_memory = cfg->debugging; - ctx->debugging = cfg->debugging; - ctx->profiling = cfg->debugging; - ctx->logging = cfg->debugging; - ctx->error = NULL; - ctx->log = stderr; - create_lock(&ctx->lock); - $stms:init_fields - init_constants(ctx); - return ctx; - }|] - ) - - GC.publicDef_ "context_free" GC.InitDecl $ \s -> - ( [C.cedecl|void $id:s(struct $id:ctx* ctx);|], - [C.cedecl|void $id:s(struct $id:ctx* ctx) { - free_constants(ctx); - free_lock(&ctx->lock); - free(ctx); - }|] - ) - - GC.publicDef_ "context_sync" GC.MiscDecl $ \s -> - ( [C.cedecl|int $id:s(struct $id:ctx* ctx);|], - [C.cedecl|int $id:s(struct $id:ctx* ctx) { - (void)ctx; - return 0; - }|] - ) diff --git a/src/Futhark/CodeGen/Backends/SequentialC/Boilerplate.hs b/src/Futhark/CodeGen/Backends/SequentialC/Boilerplate.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/CodeGen/Backends/SequentialC/Boilerplate.hs @@ -0,0 +1,113 @@ +{-# LANGUAGE QuasiQuotes #-} + +module Futhark.CodeGen.Backends.SequentialC.Boilerplate (generateBoilerplate) where + +import qualified Futhark.CodeGen.Backends.GenericC as GC +import qualified Language.C.Quote.OpenCL as C + +generateBoilerplate :: GC.CompilerM op s () +generateBoilerplate = do + cfg <- GC.publicDef "context_config" GC.InitDecl $ \s -> + ( [C.cedecl|struct $id:s;|], + [C.cedecl|struct $id:s { int debugging; };|] + ) + + GC.publicDef_ "context_config_new" GC.InitDecl $ \s -> + ( [C.cedecl|struct $id:cfg* $id:s(void);|], + [C.cedecl|struct $id:cfg* $id:s(void) { + struct $id:cfg *cfg = (struct $id:cfg*) malloc(sizeof(struct $id:cfg)); + if (cfg == NULL) { + return NULL; + } + cfg->debugging = 0; + return cfg; + }|] + ) + + GC.publicDef_ "context_config_free" GC.InitDecl $ \s -> + ( [C.cedecl|void $id:s(struct $id:cfg* cfg);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg) { + free(cfg); + }|] + ) + + GC.publicDef_ "context_config_set_debugging" GC.InitDecl $ \s -> + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int detail) { + cfg->debugging = detail; + }|] + ) + + GC.publicDef_ "context_config_set_logging" GC.InitDecl $ \s -> + ( [C.cedecl|void $id:s(struct $id:cfg* cfg, int flag);|], + [C.cedecl|void $id:s(struct $id:cfg* cfg, int detail) { + /* Does nothing for this backend. */ + (void)cfg; (void)detail; + }|] + ) + + (fields, init_fields) <- GC.contextContents + + ctx <- GC.publicDef "context" GC.InitDecl $ \s -> + ( [C.cedecl|struct $id:s;|], + [C.cedecl|struct $id:s { + int detail_memory; + int debugging; + int profiling; + int logging; + typename lock_t lock; + char *error; + typename FILE *log; + int profiling_paused; + $sdecls:fields + };|] + ) + + GC.publicDef_ "context_new" GC.InitDecl $ \s -> + ( [C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg);|], + [C.cedecl|struct $id:ctx* $id:s(struct $id:cfg* cfg) { + struct $id:ctx* ctx = (struct $id:ctx*) malloc(sizeof(struct $id:ctx)); + if (ctx == NULL) { + return NULL; + } + ctx->detail_memory = cfg->debugging; + ctx->debugging = cfg->debugging; + ctx->profiling = cfg->debugging; + ctx->logging = cfg->debugging; + ctx->error = NULL; + ctx->log = stderr; + create_lock(&ctx->lock); + $stms:init_fields + init_constants(ctx); + return ctx; + }|] + ) + + GC.publicDef_ "context_free" GC.InitDecl $ \s -> + ( [C.cedecl|void $id:s(struct $id:ctx* ctx);|], + [C.cedecl|void $id:s(struct $id:ctx* ctx) { + free_constants(ctx); + free_lock(&ctx->lock); + free(ctx); + }|] + ) + + GC.publicDef_ "context_sync" GC.MiscDecl $ \s -> + ( [C.cedecl|int $id:s(struct $id:ctx* ctx);|], + [C.cedecl|int $id:s(struct $id:ctx* ctx) { + (void)ctx; + return 0; + }|] + ) + + GC.earlyDecl [C.cedecl|static const char *size_names[0];|] + GC.earlyDecl [C.cedecl|static const char *size_vars[0];|] + GC.earlyDecl [C.cedecl|static const char *size_classes[0];|] + + GC.publicDef_ "context_config_set_size" GC.InitDecl $ \s -> + ( [C.cedecl|int $id:s(struct $id:cfg* cfg, const char *size_name, size_t size_value);|], + [C.cedecl|int $id:s(struct $id:cfg* cfg, const char *size_name, size_t size_value) { + (void)cfg; (void)size_name; (void)size_value; + return 1; + }|] + ) diff --git a/src/Futhark/CodeGen/Backends/SimpleRep.hs b/src/Futhark/CodeGen/Backends/SimpleRep.hs --- a/src/Futhark/CodeGen/Backends/SimpleRep.hs +++ b/src/Futhark/CodeGen/Backends/SimpleRep.hs @@ -680,6 +680,15 @@ cFloat32Funs :: [C.Definition] cFloat32Funs = [C.cunit| + static inline typename bool $id:(funName' "isnan32")(float x) { + return isnan(x); + } + + static inline typename bool $id:(funName' "isinf32")(float x) { + return isinf(x); + } + +$esc:("#ifdef __OPENCL_VERSION__") static inline float $id:(funName' "log32")(float x) { return log(x); } @@ -760,33 +769,6 @@ return lgamma(x); } - static inline typename bool $id:(funName' "isnan32")(float x) { - return isnan(x); - } - - static inline typename bool $id:(funName' "isinf32")(float x) { - return isinf(x); - } - - static inline typename int32_t $id:(funName' "to_bits32")(float x) { - union { - float f; - typename int32_t t; - } p; - p.f = x; - return p.t; - } - - static inline float $id:(funName' "from_bits32")(typename int32_t x) { - union { - typename int32_t f; - float t; - } p; - p.f = x; - return p.t; - } - -$esc:("#ifdef __OPENCL_VERSION__") static inline float fmod32(float x, float y) { return fmod(x, y); } @@ -809,6 +791,86 @@ return fma(a,b,c); } $esc:("#else") + static inline float $id:(funName' "log32")(float x) { + return logf(x); + } + + static inline float $id:(funName' "log2_32")(float x) { + return log2f(x); + } + + static inline float $id:(funName' "log10_32")(float x) { + return log10f(x); + } + + static inline float $id:(funName' "sqrt32")(float x) { + return sqrtf(x); + } + + static inline float $id:(funName' "exp32")(float x) { + return expf(x); + } + + static inline float $id:(funName' "cos32")(float x) { + return cosf(x); + } + + static inline float $id:(funName' "sin32")(float x) { + return sinf(x); + } + + static inline float $id:(funName' "tan32")(float x) { + return tanf(x); + } + + static inline float $id:(funName' "acos32")(float x) { + return acosf(x); + } + + static inline float $id:(funName' "asin32")(float x) { + return asinf(x); + } + + static inline float $id:(funName' "atan32")(float x) { + return atanf(x); + } + + static inline float $id:(funName' "cosh32")(float x) { + return coshf(x); + } + + static inline float $id:(funName' "sinh32")(float x) { + return sinhf(x); + } + + static inline float $id:(funName' "tanh32")(float x) { + return tanhf(x); + } + + static inline float $id:(funName' "acosh32")(float x) { + return acoshf(x); + } + + static inline float $id:(funName' "asinh32")(float x) { + return asinhf(x); + } + + static inline float $id:(funName' "atanh32")(float x) { + return atanhf(x); + } + + static inline float $id:(funName' "atan2_32")(float x, float y) { + return atan2f(x,y); + } + + static inline float $id:(funName' "gamma32")(float x) { + return tgammaf(x); + } + + static inline float $id:(funName' "lgamma32")(float x) { + return lgammaf(x); + } + static inline float fmod32(float x, float y) { return fmodf(x, y); } @@ -831,6 +893,27 @@ return fmaf(a,b,c); } $esc:("#endif") + static inline typename int32_t $id:(funName' "to_bits32")(float x) { + union { + float f; + typename int32_t t; + } p; + p.f = x; + return p.t; + } + + static inline float $id:(funName' "from_bits32")(typename int32_t x) { + union { + typename int32_t f; + float t; + } p; + p.f = x; + return p.t; + } + + static inline double fsignum32(double x) { + return $id:(funName' "isnan32")(x) ? x : ((x > 0) - (x < 0)); + } |] cFloat64Funs :: [C.Definition] @@ -960,6 +1043,10 @@ static inline double fmod64(double x, double y) { return fmod(x, y); + } + + static inline double fsignum64(double x) { + return $id:(funName' "isnan64")(x) ? x : ((x > 0) - (x < 0)); } $esc:("#ifdef __OPENCL_VERSION__") 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 @@ -420,7 +420,7 @@ instance Pretty Param where ppr (ScalarParam name ptype) = ppr ptype <+> ppr name - ppr (MemParam name space) = text "mem" <> ppr space <+> ppr name + ppr (MemParam name space) = text "mem" <> ppr space <> text " " <> ppr name instance Pretty ValueDesc where ppr (ScalarValue t ept name) = @@ -613,8 +613,22 @@ declaredIn _ = mempty instance FreeIn a => FreeIn (Functions a) where - freeIn' (Functions fs) = - foldMap (freeIn' . functionBody . snd) fs + freeIn' (Functions fs) = foldMap (onFun . snd) fs + where + onFun f = + fvBind pnames $ + freeIn' (functionBody f) <> freeIn' (functionResult f <> functionArgs f) + where + pnames = + namesFromList $ map paramName $ functionInput f <> functionOutput f + +instance FreeIn ValueDesc where + freeIn' (ArrayValue mem _ _ _ dims) = freeIn' mem <> freeIn' dims + freeIn' ScalarValue {} = mempty + +instance FreeIn ExternalValue where + freeIn' (TransparentValue vd) = freeIn' vd + freeIn' (OpaqueValue _ vds) = foldMap freeIn' vds instance FreeIn a => FreeIn (Code a) where freeIn' (x :>>: y) = diff --git a/src/Futhark/CodeGen/ImpGen/Kernels.hs b/src/Futhark/CodeGen/ImpGen/Kernels.hs --- a/src/Futhark/CodeGen/ImpGen/Kernels.hs +++ b/src/Futhark/CodeGen/ImpGen/Kernels.hs @@ -50,7 +50,17 @@ openclAtomics, cudaAtomics :: AtomicBinOp (openclAtomics, cudaAtomics) = (flip lookup opencl, flip lookup cuda) where - opencl = + opencl64 = + [ (Add Int64 OverflowUndef, Imp.AtomicAdd Int64), + (SMax Int64, Imp.AtomicSMax Int64), + (SMin Int64, Imp.AtomicSMin Int64), + (UMax Int64, Imp.AtomicUMax Int64), + (UMin Int64, Imp.AtomicUMin Int64), + (And Int64, Imp.AtomicAnd Int64), + (Or Int64, Imp.AtomicOr Int64), + (Xor Int64, Imp.AtomicXor Int64) + ] + opencl32 = [ (Add Int32 OverflowUndef, Imp.AtomicAdd Int32), (SMax Int32, Imp.AtomicSMax Int32), (SMin Int32, Imp.AtomicSMin Int32), @@ -60,7 +70,12 @@ (Or Int32, Imp.AtomicOr Int32), (Xor Int32, Imp.AtomicXor Int32) ] - cuda = opencl ++ [(FAdd Float32, Imp.AtomicFAdd Float32)] + opencl = opencl32 ++ opencl64 + cuda = + opencl + ++ [ (FAdd Float32, Imp.AtomicFAdd Float32), + (FAdd Float64, Imp.AtomicFAdd Float64) + ] compileProg :: MonadFreshNames m => diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/Base.hs b/src/Futhark/CodeGen/ImpGen/Kernels/Base.hs --- a/src/Futhark/CodeGen/ImpGen/Kernels/Base.hs +++ b/src/Futhark/CodeGen/ImpGen/Kernels/Base.hs @@ -550,13 +550,13 @@ AtomicUpdate KernelsMem KernelEnv atomicUpdateLocking atomicBinOp lam | Just ops_and_ts <- splitOp lam, - all (\(_, t, _, _) -> primBitSize t == 32) ops_and_ts = + all (\(_, t, _, _) -> primBitSize t `elem` [32, 64]) ops_and_ts = primOrCas ops_and_ts $ \space arrs bucket -> - -- If the operator is a vectorised binary operator on 32-bit values, - -- we can use a particularly efficient implementation. If the - -- operator has an atomic implementation we use that, otherwise it - -- is still a binary operator which can be implemented by atomic - -- compare-and-swap if 32 bits. + -- If the operator is a vectorised binary operator on 32/64-bit + -- values, we can use a particularly efficient + -- implementation. If the operator has an atomic implementation + -- we use that, otherwise it is still a binary operator which + -- can be implemented by atomic compare-and-swap if 32/64 bits. forM_ (zip arrs ops_and_ts) $ \(a, (op, t, x, y)) -> do -- Common variables. old <- dPrim "old" t @@ -579,13 +579,13 @@ isPrim (op, _, _, _) = isJust $ atomicBinOp op --- If the operator functions purely on single 32-bit values, we can +-- If the operator functions purely on single 32/64-bit values, we can -- use an implementation based on CAS, no matter what the operator -- does. atomicUpdateLocking _ op | [Prim t] <- lambdaReturnType op, [xp, _] <- lambdaParams op, - primBitSize t == 32 = AtomicCAS $ \space [arr] bucket -> do + primBitSize t `elem` [32, 64] = AtomicCAS $ \space [arr] bucket -> do old <- dPrim "old" t atomicUpdateCAS space t arr (tvVar old) bucket (paramName xp) $ compileBody' [xp] $ lambdaBody op @@ -703,25 +703,35 @@ ( \v -> Imp.FunExp "to_bits32" [v] int32, \v -> Imp.FunExp "from_bits32" [v] t ) + FloatType Float64 -> + ( \v -> Imp.FunExp "to_bits64" [v] int64, + \v -> Imp.FunExp "from_bits64" [v] t + ) _ -> (id, id) + + int + | primBitSize t == 32 = int32 + | otherwise = int64 + sWhile (tvExp run_loop) $ do assumed <~~ Imp.var old t x <~~ Imp.var assumed t do_op - old_bits <- dPrim "old_bits" int32 + old_bits_v <- newVName "old_bits" + dPrim_ old_bits_v int + let old_bits = Imp.var old_bits_v int sOp $ Imp.Atomic space $ Imp.AtomicCmpXchg - int32 - (tvVar old_bits) + int + old_bits_v arr' bucket_offset (toBits (Imp.var assumed t)) (toBits (Imp.var x t)) - old <~~ fromBits (untyped $ tvExp old_bits) - sWhen - (isInt32 (toBits (Imp.var assumed t)) .==. tvExp old_bits) - (run_loop <-- false) + old <~~ fromBits old_bits + let won = CmpOpExp (CmpEq int) (toBits (Imp.var assumed t)) old_bits + sWhen (isBool won) (run_loop <-- false) -- | Horizontally fission a lambda that models a binary operator. splitOp :: ASTLore lore => Lambda lore -> Maybe [(BinOp, PrimType, VName, VName)] @@ -1648,6 +1658,42 @@ localOps threadOperations $ sWhen (isActive $ zip (map tvVar is_for_thread) $ map fst dims) $ copyDWIMFix (patElemName pe) (map tvExp is_for_thread) (Var what) local_is +compileGroupResult space pe (RegTileReturns dims_n_tiles what) = do + constants <- kernelConstants <$> askEnv + + let gids = map fst $ unSegSpace space + (dims, group_tiles, reg_tiles) = unzip3 dims_n_tiles + group_tiles' = map toInt64Exp group_tiles + reg_tiles' = map toInt64Exp reg_tiles + + -- Which group tile is this group responsible for? + let group_tile_is = map Imp.vi64 gids + + -- Within the group tile, which register tile is this thread + -- responsible for? + reg_tile_is <- + mapM (dPrimVE "reg_tile_i") $ + unflattenIndex group_tiles' $ sExt64 $ kernelLocalThreadId constants + + -- Compute output array slice for the register tile belonging to + -- this thread. + let regTileSliceDim (group_tile, group_tile_i) (reg_tile, reg_tile_i) = do + tile_dim_start <- + dPrimVE "tile_dim_start" $ + reg_tile * (group_tile * group_tile_i + reg_tile_i) + return $ DimSlice tile_dim_start reg_tile 1 + reg_tile_slices <- + zipWithM + regTileSliceDim + (zip group_tiles' group_tile_is) + (zip reg_tiles' reg_tile_is) + + localOps threadOperations $ + sLoopNest (Shape reg_tiles) $ \is_in_reg_tile -> do + let dest_is = fixSlice reg_tile_slices is_in_reg_tile + src_is = reg_tile_is ++ is_in_reg_tile + sWhen (foldl1 (.&&.) $ zipWith (.<.) dest_is $ map toInt64Exp dims) $ + copyDWIMFix (patElemName pe) dest_is (Var what) src_is compileGroupResult space pe (Returns _ what) = do constants <- kernelConstants <$> askEnv in_local_memory <- arrayInLocalMemory what @@ -1673,6 +1719,8 @@ PatElem KernelsMem -> KernelResult -> InKernelGen () +compileThreadResult _ _ RegTileReturns {} = + compilerLimitationS "compileThreadResult: RegTileReturns not yet handled." compileThreadResult space pe (Returns _ what) = do let is = map (Imp.vi64 . fst) $ unSegSpace space copyDWIMFix (patElemName pe) is what [] diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/SegHist.hs b/src/Futhark/CodeGen/ImpGen/Kernels/SegHist.hs --- a/src/Futhark/CodeGen/ImpGen/Kernels/SegHist.hs +++ b/src/Futhark/CodeGen/ImpGen/Kernels/SegHist.hs @@ -199,7 +199,7 @@ bodyPassage :: KernelBody KernelsMem -> Passage bodyPassage kbody - | mempty == consumedInKernelBody (aliasAnalyseKernelBody kbody) = + | mempty == consumedInKernelBody (aliasAnalyseKernelBody mempty kbody) = MayBeMultiPass | otherwise = MustBeSinglePass diff --git a/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs b/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs --- a/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs +++ b/src/Futhark/CodeGen/ImpGen/Kernels/ToOpenCL.hs @@ -424,6 +424,9 @@ if (thread_gid >= n) return; } +$esc:("#pragma OPENCL EXTENSION cl_khr_int64_base_atomics : enable") +$esc:("#pragma OPENCL EXTENSION cl_khr_int64_extended_atomics : enable") + typedef char int8_t; typedef short int16_t; typedef int int32_t; @@ -688,10 +691,49 @@ where op' = op ++ "_" ++ pretty t ++ "_" ++ atomicSpace s + doAtomicCmpXchg s t old arr ind cmp val ty = do + ind' <- GC.compileExp $ untyped $ unCount ind + cmp' <- GC.compileExp cmp + val' <- GC.compileExp val + cast <- atomicCast s ty + GC.stm [C.cstm|$id:old = $id:op(&(($ty:cast *)$id:arr)[$exp:ind'], $exp:cmp', $exp:val');|] + where + op = "atomic_cmpxchg_" ++ pretty t ++ "_" ++ atomicSpace s + doAtomicXchg s t old arr ind val ty = do + cast <- atomicCast s ty + ind' <- GC.compileExp $ untyped $ unCount ind + val' <- GC.compileExp val + GC.stm [C.cstm|$id:old = $id:op(&(($ty:cast *)$id:arr)[$exp:ind'], $exp:val');|] + where + op = "atomic_chg_" ++ pretty t ++ "_" ++ atomicSpace s + -- First the 64-bit operations. + atomicOps s (AtomicAdd Int64 old arr ind val) = + doAtomic s Int64 old arr ind val "atomic_add" [C.cty|typename int64_t|] + atomicOps s (AtomicFAdd Float64 old arr ind val) = + doAtomic s Float64 old arr ind val "atomic_fadd" [C.cty|double|] + atomicOps s (AtomicSMax Int64 old arr ind val) = + doAtomic s Int64 old arr ind val "atomic_smax" [C.cty|typename int64_t|] + atomicOps s (AtomicSMin Int64 old arr ind val) = + doAtomic s Int64 old arr ind val "atomic_smin" [C.cty|typename int64_t|] + atomicOps s (AtomicUMax Int64 old arr ind val) = + doAtomic s Int64 old arr ind val "atomic_umax" [C.cty|unsigned int64_t|] + atomicOps s (AtomicUMin Int64 old arr ind val) = + doAtomic s Int64 old arr ind val "atomic_umin" [C.cty|unsigned int64_t|] + atomicOps s (AtomicAnd Int64 old arr ind val) = + doAtomic s Int64 old arr ind val "atomic_and" [C.cty|typename int64_t|] + atomicOps s (AtomicOr Int64 old arr ind val) = + doAtomic s Int64 old arr ind val "atomic_or" [C.cty|typename int64_t|] + atomicOps s (AtomicXor Int64 old arr ind val) = + doAtomic s Int64 old arr ind val "atomic_xor" [C.cty|typename int64_t|] + atomicOps s (AtomicCmpXchg (IntType Int64) old arr ind cmp val) = + doAtomicCmpXchg s (IntType Int64) old arr ind cmp val [C.cty|typename int64_t|] + atomicOps s (AtomicXchg (IntType Int64) old arr ind val) = + doAtomicXchg s (IntType Int64) old arr ind val [C.cty|typename int64_t|] + -- atomicOps s (AtomicAdd t old arr ind val) = doAtomic s t old arr ind val "atomic_add" [C.cty|int|] - atomicOps s (AtomicFAdd t old arr ind val) = - doAtomic s t old arr ind val "atomic_fadd" [C.cty|float|] + atomicOps s (AtomicFAdd Float32 old arr ind val) = + doAtomic s Float32 old arr ind val "atomic_fadd" [C.cty|float|] atomicOps s (AtomicSMax t old arr ind val) = doAtomic s t old arr ind val "atomic_smax" [C.cty|int|] atomicOps s (AtomicSMin t old arr ind val) = @@ -706,21 +748,10 @@ doAtomic s t old arr ind val "atomic_or" [C.cty|int|] atomicOps s (AtomicXor t old arr ind val) = doAtomic s t old arr ind val "atomic_xor" [C.cty|int|] - atomicOps s (AtomicCmpXchg t old arr ind cmp val) = do - ind' <- GC.compileExp $ untyped $ unCount ind - cmp' <- GC.compileExp cmp - val' <- GC.compileExp val - cast <- atomicCast s [C.cty|int|] - GC.stm [C.cstm|$id:old = $id:op(&(($ty:cast *)$id:arr)[$exp:ind'], $exp:cmp', $exp:val');|] - where - op = "atomic_cmpxchg_" ++ pretty t ++ "_" ++ atomicSpace s - atomicOps s (AtomicXchg t old arr ind val) = do - ind' <- GC.compileExp $ untyped $ unCount ind - val' <- GC.compileExp val - cast <- atomicCast s [C.cty|int|] - GC.stm [C.cstm|$id:old = $id:op(&(($ty:cast *)$id:arr)[$exp:ind'], $exp:val');|] - where - op = "atomic_cmpxchg_" ++ pretty t ++ "_" ++ atomicSpace s + atomicOps s (AtomicCmpXchg t old arr ind cmp val) = + doAtomicCmpXchg s t old arr ind cmp val [C.cty|int|] + atomicOps s (AtomicXchg t old arr ind val) = + doAtomicXchg s t old arr ind val [C.cty|int|] cannotAllocate :: GC.Allocate KernelOp KernelState cannotAllocate _ = diff --git a/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs b/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs --- a/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs +++ b/src/Futhark/CodeGen/ImpGen/Multicore/Base.hs @@ -118,6 +118,8 @@ compilerBugS "compileThreadResult: WriteReturns unhandled." compileThreadResult _ _ TileReturns {} = compilerBugS "compileThreadResult: TileReturns unhandled." +compileThreadResult _ _ RegTileReturns {} = + compilerBugS "compileThreadResult: RegTileReturns unhandled." freeVariables :: Imp.Code -> [VName] -> [VName] freeVariables code names = 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 @@ -34,7 +34,7 @@ ppr (DeviceInfo s) = text "device_info" <> parens (ppr s) -- | A size that can be assigned a default. -data WhichSize = LockstepWidth | NumGroups | GroupSize | TileSize | Threshold +data WhichSize = LockstepWidth | NumGroups | GroupSize | TileSize | RegTileSize | Threshold -- | A heuristic for setting the default value for something. data SizeHeuristic = SizeHeuristic @@ -56,11 +56,13 @@ SizeHeuristic "" DeviceGPU NumGroups $ 4 * max_compute_units, SizeHeuristic "" DeviceGPU GroupSize 256, SizeHeuristic "" DeviceGPU TileSize 32, + SizeHeuristic "" DeviceGPU RegTileSize 2, SizeHeuristic "" DeviceGPU Threshold $ 32 * 1024, SizeHeuristic "" DeviceCPU LockstepWidth 1, SizeHeuristic "" DeviceCPU NumGroups max_compute_units, SizeHeuristic "" DeviceCPU GroupSize 32, SizeHeuristic "" DeviceCPU TileSize 4, + SizeHeuristic "" DeviceCPU RegTileSize 1, SizeHeuristic "" DeviceCPU Threshold max_compute_units ] where diff --git a/src/Futhark/Compiler.hs b/src/Futhark/Compiler.hs --- a/src/Futhark/Compiler.hs +++ b/src/Futhark/Compiler.hs @@ -46,7 +46,9 @@ -- | If True, ignore @unsafe@. futharkSafe :: Bool, -- | Additional functions that should be exposed as entry points. - futharkEntryPoints :: [Name] + futharkEntryPoints :: [Name], + -- | If false, disable type-checking + futharkTypeCheck :: Bool } -- | The default compiler configuration. @@ -57,7 +59,8 @@ futharkWarn = True, futharkWerror = False, futharkSafe = False, - futharkEntryPoints = [] + futharkEntryPoints = [], + futharkTypeCheck = True } -- | Print a compiler error to stdout. The 'FutharkConfig' controls @@ -140,7 +143,7 @@ pipeline_config = PipelineConfig { pipelineVerbose = fst (futharkVerbose config) > NotVerbose, - pipelineValidate = True + pipelineValidate = futharkTypeCheck config } typeCheckInternalProgram :: I.Prog I.SOACS -> FutharkM () 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 @@ -2,6 +2,7 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TupleSections #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} @@ -41,7 +42,7 @@ -- * Tracking aliases AliasesAndConsumed, trackAliases, - consumedInStms, + mkStmsAliases, ) where @@ -353,6 +354,8 @@ consumed' = consumed `namesSubtract` boundNames in (map AliasDec aliases', AliasDec consumed') +-- | The aliases of the result and everything consumed in the given +-- statements. mkStmsAliases :: Aliased lore => Stms lore -> @@ -371,11 +374,6 @@ where look k = M.findWithDefault mempty k aliasmap --- | Everything consumed in the given statements and result (even --- transitively). -consumedInStms :: Aliased lore => Stms lore -> Names -consumedInStms = snd . flip mkStmsAliases [] - type AliasesAndConsumed = ( M.Map VName Names, Names @@ -386,13 +384,17 @@ AliasesAndConsumed -> Stm lore -> AliasesAndConsumed -trackAliases (aliasmap, consumed) bnd = - let pat = stmPattern bnd - als = - M.fromList $ - zip (patternNames pat) (map addAliasesOfAliases $ patternAliases pat) - aliasmap' = als <> aliasmap - consumed' = consumed <> addAliasesOfAliases (consumedInStm bnd) +trackAliases (aliasmap, consumed) stm = + let pat = stmPattern stm + pe_als = + zip (patternNames pat) $ map addAliasesOfAliases $ patternAliases pat + als = M.fromList pe_als + rev_als = foldMap revAls pe_als + revAls (v, v_als) = + M.fromList $ map (,oneName v) $ namesToList v_als + comb = M.unionWith (<>) + aliasmap' = rev_als `comb` als `comb` aliasmap + consumed' = consumed <> addAliasesOfAliases (consumedInStm stm) in (aliasmap', consumed') where addAliasesOfAliases names = names <> aliasesOfAliases names diff --git a/src/Futhark/IR/Kernels/Kernel.hs b/src/Futhark/IR/Kernels/Kernel.hs --- a/src/Futhark/IR/Kernels/Kernel.hs +++ b/src/Futhark/IR/Kernels/Kernel.hs @@ -321,9 +321,9 @@ instance (CanBeAliased (Op lore), CanBeAliased op, ASTLore lore) => CanBeAliased (HostOp lore op) where type OpWithAliases (HostOp lore op) = HostOp (Aliases lore) (OpWithAliases op) - addOpAliases (SegOp op) = SegOp $ addOpAliases op - addOpAliases (OtherOp op) = OtherOp $ addOpAliases op - addOpAliases (SizeOp op) = SizeOp op + addOpAliases aliases (SegOp op) = SegOp $ addOpAliases aliases op + 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 diff --git a/src/Futhark/IR/Kernels/Simplify.hs b/src/Futhark/IR/Kernels/Simplify.hs --- a/src/Futhark/IR/Kernels/Simplify.hs +++ b/src/Futhark/IR/Kernels/Simplify.hs @@ -94,10 +94,10 @@ <> ruleBook [ RuleOp redomapIotaToLoop, RuleOp SOAC.simplifyKnownIterationSOAC, - RuleOp SOAC.removeReplicateMapping - ] - [ RuleBasicOp removeUnnecessaryCopy, + RuleOp SOAC.removeReplicateMapping, RuleOp SOAC.liftIdentityMapping + ] + [ RuleBasicOp removeUnnecessaryCopy ] -- We turn reductions over (solely) iotas into do-loops, because there diff --git a/src/Futhark/IR/Kernels/Sizes.hs b/src/Futhark/IR/Kernels/Sizes.hs --- a/src/Futhark/IR/Kernels/Sizes.hs +++ b/src/Futhark/IR/Kernels/Sizes.hs @@ -41,6 +41,7 @@ | SizeGroup | SizeNumGroups | SizeTile + | SizeRegTile | -- | Likely not useful on its own, but querying the -- maximum can be handy. SizeLocalMemory @@ -55,10 +56,11 @@ With (. Sexp.sym "group") $ With (. Sexp.sym "num-groups") $ With (. Sexp.sym "tile") $ - With (. Sexp.sym "local-memory") $ - With - (. Sexp.list (Sexp.el (Sexp.sym "bespoke") >>> Sexp.el sexpIso >>> Sexp.el (iso fromIntegral fromIntegral . Sexp.int))) - End + With (. Sexp.sym "reg-tile") $ + With (. Sexp.sym "local-memory") $ + With + (. Sexp.list (Sexp.el (Sexp.sym "bespoke") >>> Sexp.el sexpIso >>> Sexp.el (iso fromIntegral fromIntegral . Sexp.int))) + End instance Pretty SizeClass where ppr (SizeThreshold path _) = text $ "threshold (" ++ unwords (map pStep path) ++ ")" @@ -68,6 +70,7 @@ ppr SizeGroup = text "group_size" ppr SizeNumGroups = text "num_groups" ppr SizeTile = text "tile_size" + ppr SizeRegTile = text "reg_tile_size" ppr SizeLocalMemory = text "local_memory" ppr (SizeBespoke k _) = ppr k diff --git a/src/Futhark/IR/MC/Op.hs b/src/Futhark/IR/MC/Op.hs --- a/src/Futhark/IR/MC/Op.hs +++ b/src/Futhark/IR/MC/Op.hs @@ -105,10 +105,10 @@ where type OpWithAliases (MCOp lore op) = MCOp (Aliases lore) (OpWithAliases op) - addOpAliases (ParOp par_op op) = - ParOp (addOpAliases <$> par_op) (addOpAliases op) - addOpAliases (OtherOp op) = - OtherOp $ addOpAliases op + addOpAliases aliases (ParOp par_op op) = + ParOp (addOpAliases aliases <$> par_op) (addOpAliases aliases op) + addOpAliases aliases (OtherOp op) = + OtherOp $ addOpAliases aliases op removeOpAliases (ParOp par_op op) = ParOp (removeOpAliases <$> par_op) (removeOpAliases op) diff --git a/src/Futhark/IR/Mem.hs b/src/Futhark/IR/Mem.hs --- a/src/Futhark/IR/Mem.hs +++ b/src/Futhark/IR/Mem.hs @@ -212,8 +212,8 @@ removeOpAliases (Alloc se space) = Alloc se space removeOpAliases (Inner k) = Inner $ removeOpAliases k - addOpAliases (Alloc se space) = Alloc se space - addOpAliases (Inner k) = Inner $ addOpAliases k + addOpAliases _ (Alloc se space) = Alloc se space + addOpAliases aliases (Inner k) = Inner $ addOpAliases aliases k instance Rename inner => Rename (MemOp inner) where rename (Alloc size space) = Alloc <$> rename size <*> pure space diff --git a/src/Futhark/IR/Pretty.hs b/src/Futhark/IR/Pretty.hs --- a/src/Futhark/IR/Pretty.hs +++ b/src/Futhark/IR/Pretty.hs @@ -51,13 +51,6 @@ PrettyLore lore where ppExpLore :: ExpDec lore -> Exp lore -> Maybe Doc - ppExpLore _ (If _ _ _ (IfDec ts _)) = - Just $ - stack $ - map (text . ("-- " ++)) $ - lines $ - pretty $ - text "Branch returns:" <+> ppTuple' ts ppExpLore _ _ = Nothing commastack :: [Doc] -> Doc @@ -74,14 +67,14 @@ ppr Noncommutative = text "noncommutative" instance Pretty Shape where - ppr = brackets . commasep . map ppr . shapeDims + ppr = mconcat . map (brackets . ppr) . shapeDims instance Pretty a => Pretty (Ext a) where ppr (Free e) = ppr e ppr (Ext x) = text "?" <> text (show x) instance Pretty ExtShape where - ppr = brackets . commasep . map ppr . shapeDims + ppr = mconcat . map (brackets . ppr) . shapeDims instance Pretty Space where ppr DefaultSpace = mempty @@ -115,7 +108,7 @@ instance Pretty Certificates where ppr (Certificates []) = empty - ppr (Certificates cs) = text "<" <> commasep (map ppr cs) <> text ">" + ppr (Certificates cs) = text "#" <> braces (commasep (map ppr cs)) instance PrettyLore lore => Pretty (Stms lore) where ppr = stack . map ppr . stmsToList @@ -139,6 +132,14 @@ stmAttrAnnots :: Stm lore -> [Doc] stmAttrAnnots = attrAnnots . stmAuxAttrs . stmAux +certAnnots :: Certificates -> [Doc] +certAnnots cs + | cs == mempty = [] + | otherwise = [ppr cs] + +stmCertAnnots :: Stm lore -> [Doc] +stmCertAnnots = certAnnots . stmAuxCerts . stmAux + instance Pretty (PatElemT dec) => Pretty (PatternT dec) where ppr pat = ppPattern (patternContextElements pat) (patternValueElements pat) @@ -146,35 +147,28 @@ ppr = ppr . fmap snd instance Pretty (PatElemT Type) where - ppr (PatElem name t) = ppr t <+> ppr name + ppr (PatElem name t) = ppr name <+> colon <+> ppr t instance Pretty (Param b) => Pretty (Param (a, b)) where ppr = ppr . fmap snd instance Pretty (Param DeclType) where - ppr (Param name t) = - ppr t - <+> ppr name + ppr (Param name t) = ppr name <+> colon <+> ppr t instance Pretty (Param Type) where - ppr (Param name t) = - ppr t - <+> ppr name + ppr (Param name t) = ppr name <+> colon <+> ppr t instance PrettyLore lore => Pretty (Stm lore) where - ppr bnd@(Let pat (StmAux cs _ dec) e) = + ppr bnd@(Let pat aux e) = stmannot $ align $ hang 2 $ text "let" <+> align (ppr pat) - <+> case (linebreak, ppExpLore dec e) of - (True, Nothing) -> equals e' - (_, Just ann) -> equals (ann e') - (False, Nothing) -> equals <+/> e' + <+> case (linebreak, ppExpLore (stmAuxDec aux) e) of + (True, Nothing) -> equals ppr e + (_, Just ann) -> equals (ann ppr e) + (False, Nothing) -> equals <+/> ppr e where - e' - | linebreak = ppr cs ppr e - | otherwise = ppr cs <> ppr e linebreak = case e of DoLoop {} -> True Op {} -> True @@ -182,10 +176,11 @@ Apply {} -> True BasicOp ArrayLit {} -> False BasicOp Assert {} -> True - _ -> cs /= mempty + _ -> False stmannot = case stmAttrAnnots bnd + <> stmCertAnnots bnd <> mapMaybe ppAnnot (patternElements $ stmPattern bnd) of [] -> id annots -> (align (stack annots) ) @@ -193,12 +188,12 @@ instance Pretty BasicOp where ppr (SubExp se) = ppr se ppr (Opaque e) = text "opaque" <> apply [ppr e] - ppr (ArrayLit [] rt) = - text "empty" <> parens (ppr rt) ppr (ArrayLit es rt) = case rt of Array {} -> brackets $ commastack $ map ppr es _ -> brackets $ commasep $ map ppr es + <+> colon + <+> text "[]" <> ppr rt ppr (BinOp bop x y) = ppr bop <> parens (ppr x <> comma <+> ppr y) ppr (CmpOp op x y) = ppr op <> parens (ppr x <> comma <+> ppr y) ppr (ConvOp conv x) = @@ -210,7 +205,7 @@ ppr v <> brackets (commasep (map ppr idxs)) ppr (Update src idxs se) = ppr src <+> text "with" <+> brackets (commasep (map ppr idxs)) - <+> text "<-" + <+> text "=" <+> ppr se ppr (Iota e x s et) = text "iota" <> et' <> apply [ppr e, ppr x, ppr s] where @@ -225,27 +220,29 @@ text "rearrange" <> apply [apply (map ppr perm), ppr e] ppr (Rotate es e) = text "rotate" <> apply [apply (map ppr es), ppr e] - ppr (Concat i x ys _) = - text "concat" <> text "@" <> ppr i <> apply (ppr x : map ppr ys) + ppr (Concat i x ys w) = + text "concat" <> text "@" <> ppr i <> apply (ppr w : ppr x : map ppr ys) ppr (Copy e) = text "copy" <> parens (ppr e) ppr (Manifest perm e) = text "manifest" <> apply [apply (map ppr perm), ppr e] ppr (Assert e msg (loc, _)) = text "assert" <> apply [ppr e, ppr msg, text $ show $ locStr loc] instance Pretty a => Pretty (ErrorMsg a) where - ppr (ErrorMsg parts) = commasep $ map p parts + ppr (ErrorMsg parts) = braces $ align $ commasep $ map p parts where p (ErrorString s) = text $ show s - p (ErrorInt32 x) = ppr x - p (ErrorInt64 x) = ppr x + p (ErrorInt32 x) = ppr x <+> colon <+> text "i32" + p (ErrorInt64 x) = ppr x <+> colon <+> text "i64" instance PrettyLore lore => Pretty (Exp lore) where - ppr (If c t f (IfDec _ ifsort)) = + ppr (If c t f (IfDec ret ifsort)) = text "if" <+> info' <+> ppr c text "then" <+> maybeNest t <+> text "else" <+> maybeNest f + <+> colon + <+> braces (commasep $ map ppr ret) where info' = case ifsort of IfNormal -> mempty @@ -256,7 +253,9 @@ | otherwise = nestedBlock "{" "}" $ ppr b ppr (BasicOp op) = ppr op ppr (Apply fname args _ (safety, _, _)) = - text (nameToString fname) <> safety' <> apply (map (align . pprArg) args) + text "apply" + <+> text (nameToString fname) <> safety' + <> apply (map (align . pprArg) args) where pprArg (arg, Consume) = text "*" <> ppr arg pprArg (arg, _) = ppr arg @@ -327,7 +326,7 @@ instance Pretty d => Pretty (DimIndex d) where ppr (DimFix i) = ppr i - ppr (DimSlice i n s) = ppr i <> text ":+" <> ppr n <> text "*" <> ppr s + ppr (DimSlice i n s) = ppr i <+> text ":+" <+> ppr n <+> text "*" <+> ppr s ppPattern :: (Pretty a, Pretty b) => [a] -> [b] -> Doc ppPattern [] bs = braces $ commasep $ map ppr bs 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 @@ -419,6 +419,8 @@ SSignum IntType | -- | Unsigned sign function: @usignum(2)@ = 1. USignum IntType + | -- | Floating-point sign function. + FSignum FloatType deriving (Eq, Ord, Show, Generic) instance SexpIso UnOp where @@ -429,9 +431,10 @@ With (. Sexp.list (Sexp.el (Sexp.sym "abs") >>> Sexp.el sexpIso)) $ With (. Sexp.list (Sexp.el (Sexp.sym "fabs") >>> Sexp.el sexpIso)) $ With (. Sexp.list (Sexp.el (Sexp.sym "ssignum") >>> Sexp.el sexpIso)) $ - With - (. Sexp.list (Sexp.el (Sexp.sym "usignum") >>> Sexp.el sexpIso)) - End + With (. Sexp.list (Sexp.el (Sexp.sym "usignum") >>> Sexp.el sexpIso)) $ + With + (. Sexp.list (Sexp.el (Sexp.sym "fsignum") >>> Sexp.el sexpIso)) + End -- | What to do in case of arithmetic overflow. Futhark's semantics -- are that overflow does wraparound, but for generated code (like @@ -635,27 +638,28 @@ ++ map FAbs [minBound .. maxBound] ++ map SSignum [minBound .. maxBound] ++ map USignum [minBound .. maxBound] + ++ map FSignum [minBound .. maxBound] -- | A list of all binary operators for all types. allBinOps :: [BinOp] allBinOps = concat - [ map (`Add` OverflowWrap) allIntTypes, + [ Add <$> allIntTypes <*> [OverflowWrap, OverflowUndef], map FAdd allFloatTypes, - map (`Sub` OverflowWrap) allIntTypes, + Sub <$> allIntTypes <*> [OverflowWrap, OverflowUndef], map FSub allFloatTypes, - map (`Mul` OverflowWrap) allIntTypes, + Mul <$> allIntTypes <*> [OverflowWrap, OverflowUndef], map FMul allFloatTypes, - map (`UDiv` Unsafe) allIntTypes, - map (`UDivUp` Unsafe) allIntTypes, - map (`SDiv` Unsafe) allIntTypes, - map (`SDivUp` Unsafe) allIntTypes, + UDiv <$> allIntTypes <*> [Unsafe, Safe], + UDivUp <$> allIntTypes <*> [Unsafe, Safe], + SDiv <$> allIntTypes <*> [Unsafe, Safe], + SDivUp <$> allIntTypes <*> [Unsafe, Safe], map FDiv allFloatTypes, map FMod allFloatTypes, - map (`UMod` Unsafe) allIntTypes, - map (`SMod` Unsafe) allIntTypes, - map (`SQuot` Unsafe) allIntTypes, - map (`SRem` Unsafe) allIntTypes, + UMod <$> allIntTypes <*> [Unsafe, Safe], + SMod <$> allIntTypes <*> [Unsafe, Safe], + SQuot <$> allIntTypes <*> [Unsafe, Safe], + SRem <$> allIntTypes <*> [Unsafe, Safe], map SMin allIntTypes, map UMin allIntTypes, map FMin allFloatTypes, @@ -683,7 +687,8 @@ map CmpSlt allIntTypes, map CmpSle allIntTypes, map FCmpLt allFloatTypes, - map FCmpLe allFloatTypes + map FCmpLe allFloatTypes, + [CmpLlt, CmpLle] ] -- | A list of all conversion operators for all types. @@ -710,6 +715,7 @@ doUnOp FAbs {} (FloatValue v) = Just $ FloatValue $ doFAbs v doUnOp SSignum {} (IntValue v) = Just $ IntValue $ doSSignum v doUnOp USignum {} (IntValue v) = Just $ IntValue $ doUSignum v +doUnOp FSignum {} (FloatValue v) = Just $ FloatValue $ doFSignum v doUnOp _ _ = Nothing -- | E.g., @~(~1) = 1@. @@ -732,6 +738,11 @@ doUSignum :: IntValue -> IntValue doUSignum v = intValue (intValueType v) $ signum $ intToWord64 v +-- | @fsignum(-2.0)@ = -1.0. +doFSignum :: FloatValue -> FloatValue +doFSignum (Float32Value v) = Float32Value $ signum v +doFSignum (Float64Value v) = Float64Value $ signum v + -- | Apply a 'BinOp' to an operand. Returns 'Nothing' if the -- application is mistyped, or outside the domain (e.g. division by -- zero). @@ -964,10 +975,8 @@ -- | Convert the former floating-point type to the latter. doFPConv :: FloatValue -> FloatType -> FloatValue -doFPConv (Float32Value v) Float32 = Float32Value v -doFPConv (Float64Value v) Float32 = Float32Value $ fromRational $ toRational v -doFPConv (Float64Value v) Float64 = Float64Value v -doFPConv (Float32Value v) Float64 = Float64Value $ fromRational $ toRational v +doFPConv v Float32 = Float32Value $ floatToFloat v +doFPConv v Float64 = Float64Value $ floatToDouble v -- | Convert a floating-point value to the nearest -- unsigned integer (rounding towards zero). @@ -1050,9 +1059,21 @@ intToInt = fromIntegral . intToInt64 floatToDouble :: FloatValue -> Double -floatToDouble (Float32Value v) = fromRational $ toRational v +floatToDouble (Float32Value v) + | isInfinite v, v > 0 = 1 / 0 + | isInfinite v, v < 0 = -1 / 0 + | isNaN v = 0 / 0 + | otherwise = fromRational $ toRational v floatToDouble (Float64Value v) = v +floatToFloat :: FloatValue -> Float +floatToFloat (Float64Value v) + | isInfinite v, v > 0 = 1 / 0 + | isInfinite v, v < 0 = -1 / 0 + | isNaN v = 0 / 0 + | otherwise = fromRational $ toRational v +floatToFloat (Float32Value v) = v + -- | The result type of a binary operator. binOpType :: BinOp -> PrimType binOpType (Add t _) = IntType t @@ -1108,6 +1129,7 @@ unOpType (Complement t) = IntType t unOpType (Abs t) = IntType t unOpType (FAbs t) = FloatType t +unOpType (FSignum t) = FloatType t -- | The input and output types of a conversion operator. convOpType :: ConvOp -> (PrimType, PrimType) @@ -1655,6 +1677,7 @@ ppr (FAbs t) = taggedF "fabs" t ppr (SSignum t) = taggedI "ssignum" t ppr (USignum t) = taggedI "usignum" t + ppr (FSignum t) = taggedF "fsignum" t ppr (Complement t) = taggedI "complement" t -- | The human-readable name for a 'ConvOp'. This is used to expose 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 @@ -25,6 +25,7 @@ consumedByLambda, -- * Extensibility + AliasTable, AliasedOp (..), CanBeAliased (..), ) @@ -188,6 +189,10 @@ opAliases () = [] consumedInOp () = mempty +-- | Pre-existing aliases for variables. Used to add transitive +-- aliases. +type AliasTable = M.Map VName Names + -- | The class of operations that can be given aliasing information. -- This is a somewhat subtle concept that is only used in the -- simplifier and when using "lore adapters". @@ -199,9 +204,9 @@ removeOpAliases :: OpWithAliases op -> op -- | Add aliases to this op. - addOpAliases :: op -> OpWithAliases op + addOpAliases :: AliasTable -> op -> OpWithAliases op instance CanBeAliased () where type OpWithAliases () = () removeOpAliases = id - addOpAliases = id + addOpAliases = const id 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 @@ -22,7 +22,6 @@ singleReduce, -- * Utility - getStreamAccums, scremaType, soacType, typeCheckSOAC, @@ -78,7 +77,7 @@ -- | A second-order array combinator (SOAC). data SOAC lore - = Stream SubExp (StreamForm lore) (Lambda lore) [VName] + = Stream SubExp (StreamForm lore) (Lambda lore) [SubExp] [VName] | -- | @Scatter @ -- -- @ -- -- The first SubExp is the length of the input arrays. The first @@ -109,7 +113,7 @@ instance Decorations lore => SexpIso (SOAC lore) where sexpIso = match $ - With (. Sexp.list (Sexp.el (Sexp.sym "stream") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With (. Sexp.list (Sexp.el (Sexp.sym "stream") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ With (. Sexp.list (Sexp.el (Sexp.sym "scatter") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ With (. Sexp.list (Sexp.el (Sexp.sym "hist") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ With @@ -156,8 +160,8 @@ -- | What kind of stream is this? data StreamForm lore - = Parallel StreamOrd Commutativity (Lambda lore) [SubExp] - | Sequential [SubExp] + = Parallel StreamOrd Commutativity (Lambda lore) + | Sequential deriving (Eq, Ord, Show, Generic) instance Decorations lore => SexpIso (StreamForm lore) where @@ -170,16 +174,10 @@ >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso - >>> Sexp.rest sexpIso ) ) $ With - ( . - Sexp.list - ( Sexp.el (Sexp.sym "sequential") - >>> Sexp.rest sexpIso - ) - ) + (. Sexp.list (Sexp.el (Sexp.sym "sequential"))) End -- | The essential parts of a 'Screma' factored out (everything @@ -401,18 +399,17 @@ SOACMapper flore tlore m -> SOAC flore -> m (SOAC tlore) -mapSOACM tv (Stream size form lam arrs) = +mapSOACM tv (Stream size form lam accs arrs) = Stream <$> mapOnSOACSubExp tv size <*> mapOnStreamForm form <*> mapOnSOACLambda tv lam + <*> mapM (mapOnSOACSubExp tv) accs <*> mapM (mapOnSOACVName tv) arrs where - mapOnStreamForm (Parallel o comm lam0 acc) = - Parallel o comm - <$> mapOnSOACLambda tv lam0 - <*> mapM (mapOnSOACSubExp tv) acc - mapOnStreamForm (Sequential acc) = - Sequential <$> mapM (mapOnSOACSubExp tv) acc + mapOnStreamForm (Parallel o comm lam0) = + Parallel o comm <$> mapOnSOACLambda tv lam0 + mapOnStreamForm Sequential = + pure Sequential mapSOACM tv (Scatter len lam ivs as) = Scatter <$> mapOnSOACSubExp tv len @@ -420,7 +417,7 @@ <*> mapM (mapOnSOACVName tv) ivs <*> mapM ( \(aw, an, a) -> - (,,) <$> mapOnSOACSubExp tv aw + (,,) <$> mapM (mapOnSOACSubExp tv) aw <*> pure an <*> mapOnSOACVName tv a ) @@ -487,22 +484,17 @@ -- | The type of a SOAC. soacType :: SOAC lore -> [Type] -soacType (Stream outersize form lam _) = +soacType (Stream outersize _ lam accs _) = map (substNamesInType substs) rtp where nms = map paramName $ take (1 + length accs) params substs = M.fromList $ zip nms (outersize : accs) Lambda params _ rtp = lam - accs = case form of - Parallel _ _ _ acc -> acc - Sequential acc -> acc soacType (Scatter _w lam _ivs as) = - zipWith arrayOfRow val_ts ws + zipWith arrayOfShape val_ts ws where - val_ts = - concatMap (take 1) $ - chunks ns $ - drop (sum ns) $ lambdaReturnType lam + indexes = sum $ zipWith (*) ns $ map length ws + val_ts = drop indexes $ lambdaReturnType lam (ws, ns, _) = unzip3 as soacType (Hist _len ops _bucket_fun _imgs) = do op <- ops @@ -523,21 +515,19 @@ where consumedArray v = fromMaybe v $ lookup v params_to_arrs params_to_arrs = zip (map paramName $ lambdaParams map_lam) arrs - consumedInOp (Stream _ form lam arrs) = + consumedInOp (Stream _ form lam accs arrs) = namesFromList $ subExpVars $ case form of - Sequential accs -> - map (consumedArray accs) $ namesToList $ consumedByLambda lam - Parallel _ _ _ accs -> - map (consumedArray accs) $ namesToList $ consumedByLambda lam + Sequential -> + map consumedArray $ namesToList $ consumedByLambda lam + Parallel {} -> + map consumedArray $ namesToList $ consumedByLambda lam where - consumedArray accs v = fromMaybe (Var v) $ lookup v $ paramsToInput accs + consumedArray v = fromMaybe (Var v) $ lookup v paramsToInput -- Drop the chunk parameter, which cannot alias anything. - paramsToInput accs = - zip - (map paramName $ drop 1 $ lambdaParams lam) - (accs ++ map Var arrs) + paramsToInput = + zip (map paramName $ drop 1 $ lambdaParams lam) (accs ++ map Var arrs) consumedInOp (Scatter _ _ _ as) = namesFromList $ map (\(_, _, a) -> a) as consumedInOp (Hist _ ops _ _) = @@ -559,36 +549,37 @@ where type OpWithAliases (SOAC lore) = SOAC (Aliases lore) - addOpAliases (Stream size form lam arr) = + addOpAliases aliases (Stream size form lam accs arr) = Stream size (analyseStreamForm form) - (Alias.analyseLambda lam) + (Alias.analyseLambda aliases lam) + accs arr where - analyseStreamForm (Parallel o comm lam0 acc) = - Parallel o comm (Alias.analyseLambda lam0) acc - analyseStreamForm (Sequential acc) = Sequential acc - addOpAliases (Scatter len lam ivs as) = - Scatter len (Alias.analyseLambda lam) ivs as - addOpAliases (Hist len ops bucket_fun imgs) = + analyseStreamForm (Parallel o comm lam0) = + Parallel o comm (Alias.analyseLambda aliases lam0) + analyseStreamForm Sequential = Sequential + addOpAliases aliases (Scatter len lam ivs as) = + Scatter len (Alias.analyseLambda aliases lam) ivs as + addOpAliases aliases (Hist len ops bucket_fun imgs) = Hist len - (map (mapHistOp Alias.analyseLambda) ops) - (Alias.analyseLambda bucket_fun) + (map (mapHistOp (Alias.analyseLambda aliases)) ops) + (Alias.analyseLambda aliases bucket_fun) imgs - addOpAliases (Screma w (ScremaForm scans reds map_lam) arrs) = + addOpAliases aliases (Screma w (ScremaForm scans reds map_lam) arrs) = Screma w ( ScremaForm (map onScan scans) (map onRed reds) - (Alias.analyseLambda map_lam) + (Alias.analyseLambda aliases map_lam) ) arrs where - onRed red = red {redLambda = Alias.analyseLambda $ redLambda red} - onScan scan = scan {scanLambda = Alias.analyseLambda $ scanLambda scan} + onRed red = red {redLambda = Alias.analyseLambda aliases $ redLambda red} + onScan scan = scan {scanLambda = Alias.analyseLambda aliases $ scanLambda scan} removeOpAliases = runIdentity . mapSOACM remove where @@ -657,8 +648,7 @@ -- | Type-check a SOAC. typeCheckSOAC :: TC.Checkable lore => SOAC (Aliases lore) -> TC.TypeM lore () -typeCheckSOAC (Stream size form lam arrexps) = do - let accexps = getStreamAccums form +typeCheckSOAC (Stream size form lam accexps arrexps) = do TC.require [Prim int64] size accargs <- mapM TC.checkArg accexps arrargs <- mapM lookupType arrexps @@ -673,7 +663,7 @@ TC.bad $ TC.TypeError "Stream with inconsistent accumulator type in lambda." -- check reduce's lambda, if any _ <- case form of - Parallel _ _ lam0 _ -> do + Parallel _ _ lam0 -> do let acct = map TC.argType accargs outerRetType = lambdaReturnType lam0 TC.checkLambda lam0 $ map TC.noArgAliases $ accargs ++ accargs @@ -684,7 +674,7 @@ ++ ", but stream's reduce lambda returns type " ++ prettyTuple outerRetType ++ "." - _ -> return () + Sequential -> return () -- just get the dflow of lambda on the fakearg, which does not alias -- arr, so we can later check that aliases of arr are not used inside lam. let fake_lamarrs' = map asArg lamarrs' @@ -693,10 +683,10 @@ -- Requirements: -- -- 0. @lambdaReturnType@ of @lam@ must be a list - -- [index types..., value types]. + -- [index types..., value types, ...]. -- - -- 1. The number of index types must be equal to the number of value types - -- and the number of writes to arrays in @as@. + -- 1. The number of index types and value types must be equal to the number + -- of return values from @lam@. -- -- 2. Each index type must have the type i64. -- @@ -715,15 +705,15 @@ TC.require [Prim int64] w -- 0. - let (_as_ws, as_ns, _as_vs) = unzip3 as + let (as_ws, as_ns, _as_vs) = unzip3 as + indexes = sum $ zipWith (*) as_ns $ map length as_ws rts = lambdaReturnType lam - rtsLen = length rts `div` 2 - rtsI = take rtsLen rts - rtsV = drop rtsLen rts + rtsI = take indexes rts + rtsV = drop indexes rts -- 1. - unless (rtsLen == sum as_ns) $ - TC.bad $ TC.TypeError "Scatter: Uneven number of index types, value types, and arrays outputs." + unless (length rts == sum as_ns + sum (zipWith (*) as_ns $ map length as_ws)) $ + TC.bad $ TC.TypeError "Scatter: number of index types, value types and array outputs do not match." -- 2. forM_ rtsI $ \rtI -> @@ -732,10 +722,10 @@ forM_ (zip (chunks as_ns rtsV) as) $ \(rtVs, (aw, _, a)) -> do -- All lengths must have type i64. - TC.require [Prim int64] aw + mapM_ (TC.require [Prim int64]) aw -- 3. - forM_ rtVs $ \rtV -> TC.requireI [rtV `arrayOfRow` aw] a + forM_ rtVs $ \rtV -> TC.requireI [arrayOfShape rtV aw] a -- 4. TC.consume =<< TC.lookupAliases a @@ -827,13 +817,8 @@ "Map function return type " ++ prettyTuple map_lam_ts ++ " wrong for given scan and reduction functions." --- | Get Stream's accumulators as a sub-expression list -getStreamAccums :: StreamForm lore -> [SubExp] -getStreamAccums (Parallel _ _ _ accs) = accs -getStreamAccums (Sequential accs) = accs - instance OpMetrics (Op lore) => OpMetrics (SOAC lore) where - opMetrics (Stream _ _ lam _) = + opMetrics (Stream _ _ lam _ _) = inside "Stream" $ lambdaMetrics lam opMetrics (Scatter _len lam _ivs _as) = inside "Scatter" $ lambdaMetrics lam @@ -846,26 +831,27 @@ lambdaMetrics map_lam instance PrettyLore lore => PP.Pretty (SOAC lore) where - ppr (Stream size form lam arrs) = + ppr (Stream size form lam acc arrs) = case form of - Parallel o comm lam0 acc -> + Parallel o comm lam0 -> let ord_str = if o == Disorder then "Per" else "" comm_str = case comm of Commutative -> "Comm" Noncommutative -> "" in text ("streamPar" ++ ord_str ++ comm_str) <> parens - ( ppr size <> comma ppr lam0 comma ppr lam + ( ppr size <> comma ppr lam0 <> comma + ppr lam <> comma commasep (PP.braces (commasep $ map ppr acc) : map ppr arrs) ) - Sequential acc -> + Sequential -> text "streamSeq" <> parens ( ppr size <> comma ppr lam <> comma commasep (PP.braces (commasep $ map ppr acc) : map ppr arrs) ) ppr (Scatter len lam ivs as) = - ppSOAC "scatter" len [lam] (Just (map Var ivs)) (map (\(_, n, a) -> (n, a)) as) + ppSOAC "scatter" len [lam] (Just (map Var ivs)) as ppr (Hist len ops bucket_fun imgs) = ppHist len ops bucket_fun imgs ppr (Screma w (ScremaForm scans reds map_lam) arrs) 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 @@ -40,11 +40,11 @@ import Futhark.IR.SOACS import Futhark.MonadFreshNames import qualified Futhark.Optimise.Simplify as Simplify -import Futhark.Optimise.Simplify.ClosedForm import qualified Futhark.Optimise.Simplify.Engine as Engine import Futhark.Optimise.Simplify.Lore import Futhark.Optimise.Simplify.Rule import Futhark.Optimise.Simplify.Rules +import Futhark.Optimise.Simplify.Rules.ClosedForm import Futhark.Pass import Futhark.Tools import Futhark.Transform.Rename @@ -95,20 +95,19 @@ simplifySOAC :: Simplify.SimplifiableLore lore => Simplify.SimplifyOp lore (SOAC lore) -simplifySOAC (Stream outerdim form lam arr) = do +simplifySOAC (Stream outerdim form lam nes arr) = do outerdim' <- Engine.simplify outerdim (form', form_hoisted) <- simplifyStreamForm form + nes' <- mapM Engine.simplify nes arr' <- mapM Engine.simplify arr (lam', lam_hoisted) <- Engine.simplifyLambda lam - return (Stream outerdim' form' lam' arr', form_hoisted <> lam_hoisted) + return (Stream outerdim' form' lam' nes' arr', form_hoisted <> lam_hoisted) where - simplifyStreamForm (Parallel o comm lam0 acc) = do - acc' <- mapM Engine.simplify acc + simplifyStreamForm (Parallel o comm lam0) = do (lam0', hoisted) <- Engine.simplifyLambda lam0 - return (Parallel o comm lam0' acc', hoisted) - simplifyStreamForm (Sequential acc) = do - acc' <- mapM Engine.simplify acc - return (Sequential acc', mempty) + return (Parallel o comm lam0', hoisted) + simplifyStreamForm Sequential = + return (Sequential, mempty) simplifySOAC (Scatter len lam ivs as) = do len' <- Engine.simplify len (lam', hoisted) <- Engine.simplifyLambda lam @@ -209,6 +208,8 @@ RuleOp removeUnusedSOACInput, RuleOp simplifyClosedFormReduce, RuleOp simplifyKnownIterationSOAC, + RuleOp liftIdentityMapping, + RuleOp removeDuplicateMapOutput, RuleOp fuseConcatScatter, RuleOp simplifyMapIota, RuleOp moveTransformToInput @@ -220,9 +221,7 @@ RuleOp removeDeadReduction, RuleOp removeDeadWrite, RuleBasicOp removeUnnecessaryCopy, - RuleOp liftIdentityMapping, RuleOp liftIdentityStreaming, - RuleOp removeDuplicateMapOutput, RuleOp mapOpToOp ] @@ -256,8 +255,8 @@ liftIdentityMapping :: forall lore. (Bindable lore, Simplify.SimplifiableLore lore, HasSOAC (Wise lore)) => - BottomUpRuleOp (Wise lore) -liftIdentityMapping (_, usages) pat aux op + TopDownRuleOp (Wise lore) +liftIdentityMapping _ pat aux op | Just (Screma w form arrs :: SOAC (Wise lore)) <- asSOAC op, Just fun <- isMapSOAC form = do let inputMap = M.fromList $ zip (map paramName $ lambdaParams fun) arrs @@ -270,16 +269,10 @@ checkInvariance (outId, Var v, _) (invariant, mapresult, rettype') | Just inp <- M.lookup v inputMap = - let e - | patElemName outId `UT.isConsumed` usages - || inp `UT.isConsumed` usages = - Copy inp - | otherwise = - SubExp $ Var inp - in ( (Pattern [] [outId], BasicOp e) : invariant, - mapresult, - rettype' - ) + ( (Pattern [] [outId], BasicOp (Copy inp)) : invariant, + mapresult, + rettype' + ) checkInvariance (outId, e, t) (invariant, mapresult, rettype') | freeOrConst e = ( (Pattern [] [outId], BasicOp $ Replicate (Shape [w]) e) : invariant, @@ -308,7 +301,7 @@ liftIdentityMapping _ _ _ _ = Skip liftIdentityStreaming :: BottomUpRuleOp (Wise SOACS) -liftIdentityStreaming _ (Pattern [] pes) aux (Stream w form lam arrs) +liftIdentityStreaming _ (Pattern [] pes) aux (Stream w form lam nes arrs) | (variant_map, invariant_map) <- partitionEithers $ map isInvariantRes $ zip3 map_ts map_pes map_res, not $ null invariant_map = Simplify $ do @@ -324,9 +317,9 @@ auxing aux $ letBind (Pattern [] $ fold_pes ++ variant_map_pes) $ - Op $ Stream w form lam' arrs + Op $ Stream w form lam' nes arrs where - num_folds = length $ getStreamAccums form + num_folds = length nes (fold_pes, map_pes) = splitAt num_folds pes (fold_ts, map_ts) = splitAt num_folds $ lambdaReturnType lam lam_res = bodyResult $ lambdaBody lam @@ -434,8 +427,8 @@ else Skip removeDeadMapping _ _ _ _ = Skip -removeDuplicateMapOutput :: BottomUpRuleOp (Wise SOACS) -removeDuplicateMapOutput (_, used) pat aux (Screma w form arrs) +removeDuplicateMapOutput :: TopDownRuleOp (Wise SOACS) +removeDuplicateMapOutput _ pat aux (Screma w form arrs) | Just fun <- isMapSOAC form = let ses = bodyResult $ lambdaBody fun ts = lambdaReturnType fun @@ -455,9 +448,7 @@ } auxing aux $ letBind pat' $ Op $ Screma w (mapSOAC fun') arrs forM_ copies $ \(from, to) -> - if UT.isConsumed (patElemName to) used - then letBind (Pattern [] [to]) $ BasicOp $ Copy $ patElemName from - else letBind (Pattern [] [to]) $ BasicOp $ SubExp $ Var $ patElemName from + letBind (Pattern [] [to]) $ BasicOp $ Copy $ patElemName from where checkForDuplicates (ses_ts_pes', copies) (se, t, pe) | Just (_, _, pe') <- find (\(x, _, _) -> x == se) ses_ts_pes' = @@ -697,10 +688,9 @@ zipWithM_ bindResult red_pes red_res zipWithM_ bindArrayResult map_pes map_res simplifyKnownIterationSOAC _ pat _ op - | Just (Stream (Constant k) form fold_lam arrs) <- asSOAC op, + | Just (Stream (Constant k) _ fold_lam nes arrs) <- asSOAC op, oneIsh k = Simplify $ do - let nes = getStreamAccums form - (chunk_param, acc_params, slice_params) = + let (chunk_param, acc_params, slice_params) = partitionChunkedFoldParameters (length nes) (lambdaParams fold_lam) letBindNames [paramName chunk_param] $ 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 @@ -97,7 +97,7 @@ ) import qualified Futhark.Util.Pretty as PP import GHC.Generics (Generic) -import Language.SexpGrammar as Sexp +import Language.SexpGrammar as Sexp hiding (expected) import Language.SexpGrammar.Generic import Prelude hiding (id, (.)) @@ -273,6 +273,14 @@ VName -- Tile written by this worker. -- The TileReturns must not expect more than one -- result to be written per physical thread. + | RegTileReturns + -- For each dim of result: + [ ( SubExp, -- size of this dim. + SubExp, -- block tile size for this dim. + SubExp -- reg tile size for this dim. + ) + ] + VName -- Tile returned by this worker/group. deriving (Eq, Show, Ord, Generic) instance SexpIso KernelResult where @@ -281,9 +289,10 @@ With (. Sexp.list (Sexp.el (Sexp.sym "returns") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ With (. Sexp.list (Sexp.el (Sexp.sym "write-returns") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ With (. Sexp.list (Sexp.el (Sexp.sym "concat-returns") >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ - With - (. Sexp.list (Sexp.el (Sexp.sym "tile-returns") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) - End + With (. Sexp.list (Sexp.el (Sexp.sym "tile-returns") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) $ + With + (. Sexp.list (Sexp.el (Sexp.sym "reg-tile-returns") >>> Sexp.el sexpIso >>> Sexp.el sexpIso)) + End -- | Get the root t'SubExp' corresponding values for a 'KernelResult'. kernelResultSubExp :: KernelResult -> SubExp @@ -291,6 +300,7 @@ kernelResultSubExp (WriteReturns _ arr _) = Var arr kernelResultSubExp (ConcatReturns _ _ _ v) = Var v kernelResultSubExp (TileReturns _ v) = Var v +kernelResultSubExp (RegTileReturns _ v) = Var v instance FreeIn KernelResult where freeIn' (Returns _ what) = freeIn' what @@ -299,6 +309,8 @@ freeIn' o <> freeIn' w <> freeIn' per_thread_elems <> freeIn' v freeIn' (TileReturns dims v) = freeIn' dims <> freeIn' v + freeIn' (RegTileReturns dims_n_tiles v) = + freeIn' dims_n_tiles <> freeIn' v instance ASTLore lore => FreeIn (KernelBody lore) where freeIn' (KernelBody dec stms res) = @@ -329,6 +341,10 @@ (substituteNames subst v) substituteNames subst (TileReturns dims v) = TileReturns (substituteNames subst dims) (substituteNames subst v) + substituteNames subst (RegTileReturns dims_n_tiles v) = + RegTileReturns + (substituteNames subst dims_n_tiles) + (substituteNames subst v) instance ASTLore lore => Rename (KernelBody lore) where rename (KernelBody dec stms res) = do @@ -344,10 +360,11 @@ ( ASTLore lore, CanBeAliased (Op lore) ) => + AliasTable -> KernelBody lore -> KernelBody (Aliases lore) -aliasAnalyseKernelBody (KernelBody dec stms res) = - let Body dec' stms' _ = Alias.analyseBody mempty $ Body dec stms [] +aliasAnalyseKernelBody aliases (KernelBody dec stms res) = + let Body dec' stms' _ = Alias.analyseBody aliases $ Body dec stms [] in KernelBody dec' stms' res removeKernelBodyAliases :: @@ -429,7 +446,23 @@ vt <- lookupType v unless (vt == t `arrayOfShape` Shape (map snd dims)) $ TC.bad $ TC.TypeError $ "Invalid type for TileReturns " ++ pretty v + checkKernelResult (RegTileReturns dims_n_tiles arr) t = do + mapM_ (TC.require [Prim int64]) dims + mapM_ (TC.require [Prim int64]) blk_tiles + mapM_ (TC.require [Prim int64]) reg_tiles + -- assert that arr is of element type t and shape (rev outer_tiles ++ reg_tiles) + arr_t <- lookupType arr + unless (arr_t == expected) $ + TC.bad . TC.TypeError $ + "Invalid type for TileReturns. Expected:\n " + ++ pretty expected + ++ ",\ngot:\n " + ++ pretty arr_t + where + (dims, blk_tiles, reg_tiles) = unzip3 dims_n_tiles + expected = t `arrayOfShape` Shape (blk_tiles ++ reg_tiles) + kernelBodyMetrics :: OpMetrics (Op lore) => KernelBody lore -> MetricsM () kernelBodyMetrics = mapM_ stmMetrics . kernelBodyStms @@ -459,10 +492,14 @@ SplitContiguous -> mempty SplitStrided stride -> text "Strided" <> parens (ppr stride) ppr (TileReturns dims v) = - text "tile" - <> parens (commasep $ map onDim dims) <+> ppr v + "tile" <> parens (commasep $ map onDim dims) <+> ppr v where - onDim (dim, tile) = ppr dim <+> text "/" <+> ppr tile + onDim (dim, tile) = ppr dim <+> "/" <+> ppr tile + ppr (RegTileReturns dims_n_tiles v) = + "blkreg_tile" <> parens (commasep $ map onDim dims_n_tiles) <+> ppr v + where + onDim (dim, blk_tile, reg_tile) = + ppr dim <+> "/" <+> parens (ppr blk_tile <+> "*" <+> ppr reg_tile) -- | Do we need group-virtualisation when generating code for the -- segmented operation? In most cases, we do, but for some simple @@ -573,6 +610,8 @@ t `arrayOfRow` w segResultShape _ t (TileReturns dims _) = t `arrayOfShape` Shape (map fst dims) +segResultShape _ t (RegTileReturns dims_n_tiles _) = + t `arrayOfShape` Shape (map (\(dim, _, _) -> dim) dims_n_tiles) -- | The return type of a 'SegOp'. segOpType :: SegOp lvl lore -> [Type] @@ -954,13 +993,13 @@ where type OpWithAliases (SegOp lvl lore) = SegOp lvl (Aliases lore) - addOpAliases = runIdentity . mapSegOpM alias + addOpAliases aliases = runIdentity . mapSegOpM alias where alias = SegOpMapper return - (return . Alias.analyseLambda) - (return . aliasAnalyseKernelBody) + (return . Alias.analyseLambda aliases) + (return . aliasAnalyseKernelBody aliases) return return @@ -1066,6 +1105,10 @@ <*> Engine.simplify what simplify (TileReturns dims what) = TileReturns <$> Engine.simplify dims <*> Engine.simplify what + simplify (RegTileReturns dims_n_tiles what) = + RegTileReturns + <$> Engine.simplify dims_n_tiles + <*> Engine.simplify what mkWiseKernelBody :: (ASTLore lore, CanBeWise (Op lore)) => @@ -1464,10 +1507,10 @@ SegOp lvl lore -> m [ExpReturns] segOpReturns k@(SegMap _ _ _ kbody) = - kernelBodyReturns kbody =<< (extReturns <$> opType k) + kernelBodyReturns kbody . extReturns =<< opType k segOpReturns k@(SegRed _ _ _ _ kbody) = - kernelBodyReturns kbody =<< (extReturns <$> opType k) + kernelBodyReturns kbody . extReturns =<< opType k segOpReturns k@(SegScan _ _ _ _ kbody) = - kernelBodyReturns kbody =<< (extReturns <$> opType k) + kernelBodyReturns kbody . extReturns =<< opType k segOpReturns (SegHist _ _ ops _ _) = concat <$> mapM (mapM varReturns . histDest) ops 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 @@ -416,7 +416,6 @@ Opaque SubExp | -- | Array literals, e.g., @[ [1+x, 3], [2, 1+4] ]@. -- Second arg is the element type of the rows of the array. - -- Scalar operations ArrayLit [SubExp] Type | -- | Unary operation. UnOp UnOp SubExp 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 @@ -379,8 +379,7 @@ -- | How to index a single dimension of an array. data DimIndex d = -- | Fix index in this dimension. - DimFix - d + DimFix d | -- | @DimSlice start_offset num_elems stride@. DimSlice d d d deriving (Eq, Ord, Show, Generic) diff --git a/src/Futhark/Internalise.hs b/src/Futhark/Internalise.hs --- a/src/Futhark/Internalise.hs +++ b/src/Futhark/Internalise.hs @@ -1173,8 +1173,8 @@ arrs <- internaliseExpToVars "stream_input" arr lam' <- internaliseStreamMapLambda internaliseLambda lam $ map I.Var arrs w <- arraysSize 0 <$> mapM lookupType arrs - let form = I.Parallel o Commutative (I.Lambda [] (mkBody mempty []) []) [] - letTupExp' desc $ I.Op $ I.Stream w form lam' arrs + let form = I.Parallel o Commutative (I.Lambda [] (mkBody mempty []) []) + letTupExp' desc $ I.Op $ I.Stream w form lam' [] arrs internaliseStreamRed :: String -> @@ -1237,7 +1237,7 @@ map (I.Var . paramName) lam_acc_params ++ lam_res' return $ resultBody new_lam_res - let form = I.Parallel o comm lam0' nes + let form = I.Parallel o comm lam0' lam' = I.Lambda { lambdaParams = lam_params', @@ -1245,7 +1245,7 @@ lambdaReturnType = nes_ts } w <- arraysSize 0 <$> mapM lookupType arrs - letTupExp' desc $ I.Op $ I.Stream w form lam' arrs + letTupExp' desc $ I.Op $ I.Stream w form lam' nes arrs internaliseExp1 :: String -> E.Exp -> InternaliseM I.SubExp internaliseExp1 desc e = do @@ -1647,7 +1647,9 @@ handleRest [x] "!" = Just $ complementF x handleRest [x] "opaque" = Just $ \desc -> mapM (letSubExp desc . BasicOp . Opaque) =<< internaliseExp "opaque_arg" x - handleRest [E.TupLit [a, si, v] _] "scatter" = Just $ scatterF a si v + handleRest [E.TupLit [a, si, v] _] "scatter" = Just $ scatterF 1 a si v + handleRest [E.TupLit [a, si, v] _] "scatter_2d" = Just $ scatterF 2 a si v + handleRest [E.TupLit [a, si, v] _] "scatter_3d" = Just $ scatterF 3 a si v handleRest [E.TupLit [n, m, arr] _] "unflatten" = Just $ \desc -> do arrs <- internaliseExpToVars "unflatten_arr" arr n' <- internaliseExp1 "n" n @@ -1764,13 +1766,12 @@ _ -> error "Futhark.Internalise.internaliseExp: non-int/bool type in Complement" - scatterF a si v desc = do - si' <- letExp "write_si" . BasicOp . SubExp =<< internaliseExp1 "write_arg_i" si + scatterF dim a si v desc = do + si' <- internaliseExpToVars "write_arg_i" si svs <- internaliseExpToVars "write_arg_v" v sas <- internaliseExpToVars "write_arg_a" a - si_shape <- I.arrayShape <$> lookupType si' - let si_w = shapeSize 0 si_shape + si_w <- I.arraysSize 0 <$> mapM lookupType si' sv_ts <- mapM lookupType svs svs' <- forM (zip svs sv_ts) $ \(sv, sv_t) -> do @@ -1793,21 +1794,21 @@ letExp (baseString sv ++ "_write_sv") $ I.BasicOp $ I.Reshape (reshapeOuter [DimCoercion si_w] 1 sv_shape) sv - indexType <- rowType <$> lookupType si' - indexName <- newVName "write_index" + indexType <- fmap rowType <$> mapM lookupType si' + indexName <- mapM (\_ -> newVName "write_index") indexType valueNames <- replicateM (length sv_ts) $ newVName "write_value" sa_ts <- mapM lookupType sas - let bodyTypes = replicate (length sv_ts) indexType ++ map rowType sa_ts - paramTypes = indexType : map rowType sv_ts - bodyNames = indexName : valueNames + let bodyTypes = concat (replicate (length sv_ts) indexType) ++ map (I.stripArray dim) sa_ts + paramTypes = indexType <> map rowType sv_ts + bodyNames = indexName <> valueNames bodyParams = zipWith I.Param bodyNames paramTypes -- This body is pretty boring right now, as every input is exactly the output. -- But it can get funky later on if fused with something else. body <- localScope (scopeOfLParams bodyParams) $ insertStmsM $ do - let outs = replicate (length valueNames) indexName ++ valueNames + let outs = concat (replicate (length valueNames) indexName) ++ valueNames results <- forM outs $ \name -> letSubExp "write_res" $ I.BasicOp $ I.SubExp $ I.Var name ensureResultShape @@ -1822,9 +1823,9 @@ I.lambdaReturnType = bodyTypes, I.lambdaBody = body } - sivs = si' : svs' + sivs = si' <> svs' - let sa_ws = map (arraySize 0) sa_ts + let sa_ws = map (Shape . take dim . arrayDims) sa_ts letTupExp' desc $ I.Op $ I.Scatter si_w lam sivs $ zip3 sa_ws (repeat 1) sas funcall :: @@ -1990,7 +1991,7 @@ w write_lam (classes : all_offsets ++ arrs) - $ zip3 (repeat w) (repeat 1) blanks + $ zip3 (repeat $ Shape [w]) (repeat 1) blanks sizes' <- letSubExp "partition_sizes" $ I.BasicOp $ diff --git a/src/Futhark/Internalise/Defunctionalise.hs b/src/Futhark/Internalise/Defunctionalise.hs --- a/src/Futhark/Internalise/Defunctionalise.hs +++ b/src/Futhark/Internalise/Defunctionalise.hs @@ -1,6 +1,7 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE Trustworthy #-} +{-# LANGUAGE TupleSections #-} -- | Defunctionalization of typed, monomorphic Futhark programs without modules. module Futhark.Internalise.Defunctionalise (transformProg) where @@ -408,10 +409,9 @@ <> patternArraySizes pat notSize = not . (`S.member` sizes_of_arrays) (fields, env) = - second M.fromList $ - unzip $ - map closureFromDynamicFun $ - filter (notSize . fst) $ M.toList used_env + second M.fromList . unzip . map closureFromDynamicFun + . filter (notSize . fst) + $ M.toList used_env return ( RecordLit fields loc, @@ -924,8 +924,10 @@ | otherwise -> return (e', IntrinsicSV) _ -> error $ - "Application of an expression that is neither a static lambda " - ++ "nor a dynamic function, but has static value: " + "Application of an expression\n" + ++ pretty e1 + ++ "\nthat is neither a static lambda " + ++ "nor a dynamic function, but has static value:\n" ++ show sv1 defuncApply depth e@(Var qn (Info t) loc) = do let (argtypes, _) = unfoldFunType t @@ -1128,8 +1130,12 @@ -- the pattern wins out. This is important when matching a -- nonunique pattern with a unique value. if orderZeroSV sv - then M.singleton vn $ Binding Nothing $ Dynamic t - else M.singleton vn $ Binding Nothing sv + then dim_env <> M.singleton vn (Binding Nothing $ Dynamic t) + else dim_env <> M.singleton vn (Binding Nothing sv) + where + dim_env = + M.fromList $ map (,i64) $ S.toList $ typeDimNames t + i64 = Binding Nothing $ Dynamic $ Scalar $ Prim $ Signed Int64 matchPatternSV (Wildcard _ _) _ = mempty matchPatternSV (PatternAscription pat _ _) sv = matchPatternSV pat sv matchPatternSV PatternLit {} _ = mempty 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 @@ -190,36 +190,40 @@ ModMod _ -> error $ "Cannot apply non-parametric module at " ++ locStr loc ModFun f_abs f_closure f_p f_body -> - bindingAbs (f_abs <> S.fromList (unInfo (modParamAbs f_p))) $ - extendAbsTypes b_substs $ - extendScope f_closure $ - generating $ do - outer_substs <- scopeSubsts <$> askScope - abs <- asks envAbs - let forward (k, v) = (lookupSubst k outer_substs, v) - p_substs' = M.fromList $ map forward $ M.toList p_substs - abs_substs = - M.filterWithKey (const . flip S.member abs) $ - M.map (`lookupSubst` scopeSubsts (modScope arg_mod)) p_substs' - <> scopeSubsts f_closure - <> scopeSubsts (modScope arg_mod) - extendScope - ( Scope - abs_substs - ( M.singleton (modParamName f_p) $ - substituteInMod p_substs' arg_mod - ) + bindingAbs (f_abs <> S.fromList (unInfo (modParamAbs f_p))) + . extendAbsTypes b_substs + . extendScope f_closure + . generating + $ do + outer_substs <- scopeSubsts <$> askScope + abs <- asks envAbs + let forward (k, v) = (lookupSubst k outer_substs, v) + p_substs' = M.fromList $ map forward $ M.toList p_substs + keep k _ = + k `M.member` p_substs' + || k `S.member` abs + abs_substs = + M.filterWithKey keep $ + M.map (`lookupSubst` scopeSubsts (modScope arg_mod)) p_substs' + <> scopeSubsts f_closure + <> scopeSubsts (modScope arg_mod) + extendScope + ( Scope + abs_substs + ( M.singleton (modParamName f_p) $ + substituteInMod p_substs' arg_mod ) - $ do - substs <- scopeSubsts <$> askScope - x <- evalModExp f_body - return $ - addSubsts abs abs_substs $ - -- The next one is dubious, but is necessary to - -- propagate substitutions from the argument (see - -- modules/functor24.fut). - addSubstsModMod (scopeSubsts $ modScope arg_mod) $ - substituteInMod (b_substs <> substs) x + ) + $ do + substs <- scopeSubsts <$> askScope + x <- evalModExp f_body + return $ + addSubsts abs abs_substs $ + -- The next one is dubious, but is necessary to + -- propagate substitutions from the argument (see + -- modules/functor24.fut). + addSubstsModMod (scopeSubsts $ modScope arg_mod) $ + substituteInMod (b_substs <> substs) x where addSubsts abs substs (ModFun mabs (Scope msubsts mods) mp me) = ModFun (abs <> mabs) (Scope (substs <> msubsts) mods) mp me diff --git a/src/Futhark/Internalise/FreeVars.hs b/src/Futhark/Internalise/FreeVars.hs --- a/src/Futhark/Internalise/FreeVars.hs +++ b/src/Futhark/Internalise/FreeVars.hs @@ -86,8 +86,8 @@ If e1 e2 e3 _ _ -> freeVars e1 <> freeVars e2 <> freeVars e3 Apply e1 e2 _ _ _ -> freeVars e1 <> freeVars e2 Negate e _ -> freeVars e - Lambda pats e0 _ _ _ -> - (sizes (foldMap patternDimNames pats) <> freeVars e0) + Lambda pats e0 _ (Info (_, t)) _ -> + (sizes (foldMap patternDimNames pats) <> freeVars e0 <> sizes (typeDimNames t)) `withoutM` foldMap patternVars pats OpSection {} -> mempty OpSectionLeft _ _ e _ _ _ -> freeVars e diff --git a/src/Futhark/Internalise/Monomorphise.hs b/src/Futhark/Internalise/Monomorphise.hs --- a/src/Futhark/Internalise/Monomorphise.hs +++ b/src/Futhark/Internalise/Monomorphise.hs @@ -366,7 +366,7 @@ transformExp (OpSection qn t loc) = transformExp $ Var qn t loc transformExp (OpSectionLeft fname (Info t) e arg ret loc) = do - let (Info (xtype, xargext), Info ytype) = arg + let (Info (xp, xtype, xargext), Info (yp, ytype)) = arg (Info rettype, Info retext) = ret fname' <- transformFName loc fname $ toStruct t e' <- transformExp e @@ -375,12 +375,12 @@ (Just e') Nothing t - (xtype, xargext) - (ytype, Nothing) + (xp, xtype, xargext) + (yp, ytype, Nothing) (rettype, retext) loc transformExp (OpSectionRight fname (Info t) e arg (Info rettype) loc) = do - let (Info xtype, Info (ytype, yargext)) = arg + let (Info (xp, xtype), Info (yp, ytype, yargext)) = arg fname' <- transformFName loc fname $ toStruct t e' <- transformExp e desugarBinOpSection @@ -388,8 +388,8 @@ Nothing (Just e') t - (xtype, Nothing) - (ytype, yargext) + (xp, xtype, Nothing) + (yp, ytype, yargext) (rettype, []) loc transformExp (ProjectSection fields (Info t) loc) = @@ -518,46 +518,53 @@ Maybe Exp -> Maybe Exp -> PatternType -> - (StructType, Maybe VName) -> - (StructType, Maybe VName) -> + (PName, StructType, Maybe VName) -> + (PName, StructType, Maybe VName) -> (PatternType, [VName]) -> SrcLoc -> MonoM Exp -desugarBinOpSection op e_left e_right t (xtype, xext) (ytype, yext) (rettype, retext) loc = do - (wrap_left, e1, p1) <- makeVarParam e_left $ fromStruct xtype - (wrap_right, e2, p2) <- makeVarParam e_right $ fromStruct ytype +desugarBinOpSection op e_left e_right t (xp, xtype, xext) (yp, ytype, yext) (rettype, retext) loc = do + (v1, wrap_left, e1, p1) <- makeVarParam e_left $ fromStruct xtype + (v2, wrap_right, e2, p2) <- makeVarParam e_right $ fromStruct ytype let apply_left = Apply op e1 (Info (Observe, xext)) - (Info $ foldFunType [fromStruct ytype] t, Info []) + (Info $ Scalar $ Arrow mempty yp (fromStruct ytype) t, Info []) loc + rettype' = + let onDim (NamedDim d) + | Named p <- xp, qualLeaf d == p = NamedDim $ qualName v1 + | Named p <- yp, qualLeaf d == p = NamedDim $ qualName v2 + onDim d = d + in first onDim rettype body = Apply apply_left e2 (Info (Observe, yext)) - (Info rettype, Info retext) + (Info rettype', Info retext) loc - rettype' = toStruct rettype - return $ wrap_left $ wrap_right $ Lambda (p1 ++ p2) body Nothing (Info (mempty, rettype')) loc + rettype'' = toStruct rettype' + return $ wrap_left $ wrap_right $ Lambda (p1 ++ p2) body Nothing (Info (mempty, rettype'')) loc where patAndVar argtype = do x <- newNameFromString "x" pure - ( Id x (Info argtype) mempty, + ( x, + Id x (Info argtype) mempty, Var (qualName x) (Info argtype) mempty ) makeVarParam (Just e) argtype = do - (pat, var_e) <- patAndVar argtype + (v, pat, var_e) <- patAndVar argtype let wrap body = LetPat pat e body (Info (typeOf body), Info mempty) mempty - return (wrap, var_e, []) + return (v, wrap, var_e, []) makeVarParam Nothing argtype = do - (pat, var_e) <- patAndVar argtype - return (id, var_e, [pat]) + (v, pat, var_e) <- patAndVar argtype + return (v, id, var_e, [pat]) desugarProjectSection :: [Name] -> PatternType -> SrcLoc -> MonoM Exp desugarProjectSection fields (Scalar (Arrow _ _ t1 t2)) loc = do @@ -668,17 +675,7 @@ Just (ConstDim x) -> Just $ Literal (SignedValue $ Int64Value $ fromIntegral x) mempty _ -> - Nothing - -explicitSizes :: StructType -> MonoType -> S.Set VName -explicitSizes t1 t2 = - execState (matchDims onDims t1 t2) mempty `S.intersection` mustBeExplicit t1 - where - onDims d1 d2 = do - case (d1, d2) of - (NamedDim v, MonoKnown _) -> modify $ S.insert $ qualLeaf v - _ -> return () - return d1 + Just $ Literal (SignedValue $ Int64Value 0) mempty -- Monomorphising higher-order functions can result in function types -- where the same named parameter occurs in multiple spots. When @@ -719,7 +716,7 @@ params' = map (substPattern entry substPatternType) params bind_t' = substTypesAny (`M.lookup` substs') bind_t (shape_params_explicit, shape_params_implicit) = - partition ((`S.member` explicitSizes bind_t' t) . typeParamName) $ + partition ((`S.member` mustBeExplicit bind_t') . typeParamName) $ shape_params ++ t_shape_params (params'', rrs) <- unzip <$> mapM transformPattern params' diff --git a/src/Futhark/Optimise/BlkRegTiling.hs b/src/Futhark/Optimise/BlkRegTiling.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/BlkRegTiling.hs @@ -0,0 +1,1022 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeFamilies #-} + +-- | Perform a restricted form of block+register tiling corresponding to +-- the following pattern: +-- * a redomap is quasi-perfectly nested inside a kernel with at +-- least two parallel dimension (the perfectly nested restriction +-- is relaxed a bit to allow for SGEMM); +-- * all streamed arrays of redomap are one dimensional; +-- * 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 stream's result is a tuple of scalar values, which are +-- also the "thread-in-space" return of the kernel. +-- * We have further restrictions that in principle can be relaxed: +-- the redomap has exactly two array input +-- the redomap produces one scalar result +-- the kernel produces one scalar result +module Futhark.Optimise.BlkRegTiling (mmBlkRegTiling, doRegTiling3D) where + +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.Kernels +import Futhark.MonadFreshNames +import Futhark.Optimise.TileLoops.Shared +import Futhark.Tools +import Futhark.Transform.Rename + +mmBlkRegTiling :: Stm Kernels -> TileM (Maybe (Stms Kernels, Stm Kernels)) +mmBlkRegTiling (Let pat aux (Op (SegOp (SegMap SegThread {} seg_space ts old_kbody)))) + | KernelBody () kstms [Returns ResultMaySimplify (Var res_nm)] <- old_kbody, + -- 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] <- patternValueElements 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) <- runBinder $ 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 + + gid_x <- newVName "gid_x" + gid_y <- newVName "gid_y" + gid_flat <- newVName "gid_flat" + + ---- in this binder: outer seggroup ---- + (ret_seggroup, stms_seggroup) <- runBinder $ 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'] + return [Var 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 kk0 (thd_res_merge, a_loc_init', b_loc_init') epilogue = do + kk <- letExp "kk" =<< toExp (le64 kk0 * pe64 tk) + a_loc <- forLoop ry [a_loc_init'] $ \i0 [a_loc_merge] -> do + loop_a_loc <- forLoop tk_div_tx [a_loc_merge] $ \k0 [a_loc_merge'] -> do + scatter_a_loc <- segScatter2D "A_glb2loc" a_loc_sz a_loc_merge' segthd_lvl (ty, tx) $ + \(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)]) + return (a_elem, a_loc_ind) + resultBodyM $ map Var scatter_a_loc + resultBodyM [Var loop_a_loc] + + -- copy B from global to shared memory + b_loc <- forLoop tk_div_ty [b_loc_init'] $ \k0 [b_loc_merge] -> do + loop_b_loc <- forLoop rx [b_loc_merge] $ \j0 [b_loc_merge'] -> do + scatter_b_loc <- segScatter2D + "B_glb2loc" + b_loc_sz + b_loc_merge' + segthd_lvl + (ty, tx) + $ \(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)]) + return (b_elem, b_loc_ind) + resultBodyM $ map Var scatter_b_loc + resultBodyM [Var loop_b_loc] + + -- 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] + return $ map Var [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'] -> + resultBodyM =<< letTupExp' "foo" + =<< eIf + ( toExp $ + 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_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] + return [Var css] + + resultBodyM $ map Var redomap_res + ) + (resultBodyM [Var acc_merge]) + return [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 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 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 return 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'] + return [Var rss] + let rssss : _ = rssss_list + return 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 return 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 + + return [RegTileReturns 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 + return $ Let pat aux $ Op $ SegOp $ SegMap level' space' ts kbody' + return $ Just (host_stms, new_kernel) +mmBlkRegTiling _ = return Nothing + +ceilDiv :: MonadBinder m => SubExp -> SubExp -> m (Exp (Lore m)) +ceilDiv x y = pure $ BasicOp $ BinOp (SDivUp Int64 Unsafe) x y + +scratch :: MonadBinder 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 :: MonadBinder 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 + indices = map (DimFix . Var) outer_indices ++ inner_slices + letExp se_desc $ BasicOp $ Index arr indices + +update :: MonadBinder m => String -> VName -> [VName] -> VName -> m VName +update se_desc arr indices new_elem = update' se_desc arr indices (Var new_elem) + +update' :: MonadBinder m => String -> VName -> [VName] -> SubExp -> m VName +update' se_desc arr indices new_elem = + letExp se_desc $ BasicOp $ Update arr (map (DimFix . Var) indices) new_elem + +forLoop' :: + SubExp -> -- loop var + [VName] -> -- loop inits + ( VName -> + [VName] -> -- (loop var -> loop inits -> loop body) + Binder Kernels (Body Kernels) + ) -> + Binder Kernels [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 <- + runBodyBinder . 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] -> Binder Kernels (Body Kernels)) -> + Binder Kernels VName +forLoop i_bound merge body = do + res_list <- forLoop' i_bound merge body + return $ 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 Kernels -> + [VName] -> + [VName] -> + Stms Kernels +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 lam_res -> + 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 Kernels -> + (Stms Kernels, Maybe (Stm Kernels), Stms Kernels) +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 Kernels, M.Map VName (Stm Kernels)) -> + Stm Kernels -> + Maybe (Stms Kernels, M.Map VName (Stm Kernels)) +processIndirections arrs _ acc stm@(Let patt _ (BasicOp (Index _ _))) + | Just (ss, tab) <- acc, + [p] <- patternValueElements 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 <- patternValueElements 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 -> Binder Kernels (SubExp, SubExp) +getParTiles (t_str, r_str) (t_name, r_name) len_dim = + case len_dim of + Constant (IntValue (Int64Value 8)) -> + return (se8, se1) + Constant (IntValue (Int64Value 16)) -> + return (se8, se2) + Constant (IntValue (Int64Value 32)) -> + return (se8, se4) + _ -> do + t <- letSubExp t_str $ Op $ SizeOp $ GetSize t_name SizeTile + r <- letSubExp r_str $ Op $ SizeOp $ GetSize r_name SizeRegTile + return (t, r) + +getSeqTile :: String -> Name -> SubExp -> SubExp -> SubExp -> Binder Kernels 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 Kernels -> TileM (Maybe (Stms Kernels, Stm Kernels)) +-- mmBlkRegTiling (Let pat aux (Op (SegOp (SegMap SegThread{} seg_space ts old_kbody)))) +doRegTiling3D :: Stm Kernels -> TileM (Maybe (Stms Kernels, Stm Kernels)) +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 <- patternValueElements 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, + Pattern [] _ <- pat, + all primType kertp, + all (variantToDim variance gtid_z) ker_res_nms = do + -- HERE STARTS THE IMPLEMENTATION: + (new_kernel, host_stms) <- runBinder $ 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 + + 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) <- runBinder $ 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'] + return $ Var 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) <- runBinder $ 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 + return (y_elm, y_ind) + + let ret = WriteReturns [rz] loc_Y_nm [([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 + return 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] + 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 -> return 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) -> + forM (zip kertp redomap_res) $ \(res_tp, res) -> do + rss_init <- scratch "rss_init" (elemType res_tp) [rz, se1, se1] + fmap Var $ + forLoop rz [rss_init] $ \i [rss] -> do + let slice = [DimFix $ Var i, DimFix se0, DimFix se0] + thread_res <- index "thread_res" res [ltid_y, ltid_x, i] + rss' <- letSubExp "rss" $ BasicOp $ Update rss slice $ Var thread_res + resultBodyM [rss'] + else segMap3D "rssss" segthd_lvl ResultPrivate (se1, ty, tx) $ \(_ltid_z, ltid_y, ltid_x) -> do + letBindNames [gtid_y] =<< toExp (le64 jj1 + le64 ltid_y) + letBindNames [gtid_x] =<< toExp (le64 jj2 + le64 ltid_x) + rss_init <- forM kertp $ \res_tp -> + scratch "rss_init" (elemType res_tp) [rz, se1, se1] + rss <- forLoop' rz rss_init $ \i rss_merge -> do + letBindNames [gtid_z] =<< toExp (le64 ii + le64 i) + forM_ (zip redomap_orig_res redomap_res) $ \(o_res, n_res) -> do + c <- index "redomap_thd" n_res [ltid_y, ltid_x, i] + letBindNames [patElemName o_res] =<< toExp (le64 c) + return c + res_els <- + letTupExp' "res_elem" + =<< eIf + ( toExp $ + le64 gtid_y .<. pe64 d_Ky + .&&. le64 gtid_x .<. pe64 d_Kx + .&&. le64 gtid_z .<. pe64 d_M + ) + ( do + addStms code2' + resultBodyM $ map Var ker_res_nms + ) + (eBody $ map eBlank kertp) + rss' <- forM (zip res_els rss_merge) $ \(res_el, rs_merge) -> do + let slice = [DimFix $ Var i, DimFix se0, DimFix se0] + letSubExp "rss" $ BasicOp $ Update rs_merge slice res_el + resultBodyM rss' + return $ map Var rss + + ---------------------------------------------------------------- + -- Finally, reshape the result arrays for the RegTileReturn --- + ---------------------------------------------------------------- + let regtile_ret_dims = + map (\(_, sz) -> (sz, se1, se1)) rem_outer_dims + ++ [(d_M, se1, rz), (d_Ky, ty, se1), (d_Kx, tx, se1)] + + epilogue_res' <- forM epilogue_res $ \res -> + if null rem_outer_dims + then return res + else do + -- Add dummy dimensions to tile to reflect the outer dimensions + res_tp' <- lookupType res + let (block_dims, rest_dims) = splitAt 2 $ arrayDims res_tp' + ones = map (const se1) rem_outer_dims + new_shape = concat [ones, block_dims, ones, rest_dims] + letExp "res_reshaped" $ BasicOp $ Reshape (map DimNew new_shape) res + + return $ map (RegTileReturns regtile_ret_dims) epilogue_res' + -- END (ret_seggroup, stms_seggroup) <- runBinder $ do + let level' = SegGroup (Count grid_size) (Count group_size) SegNoVirt + space' = SegSpace gid_flat (rem_outer_dims ++ [(gid_z, gridDim_z), (gid_y, gridDim_y), (gid_x, gridDim_x)]) + kbody' = KernelBody () stms_seggroup ret_seggroup + + return $ Let pat aux $ Op $ SegOp $ SegMap level' space' kertp kbody' + -- END (new_kernel, host_stms) <- runBinder $ do + return $ Just (host_stms, new_kernel) + where + getResNm (Returns ResultMaySimplify (Var res_nm)) = Just res_nm + getResNm _ = Nothing + + limitTile :: String -> SubExp -> SubExp -> Binder Kernels SubExp + limitTile t_str t d_K = letSubExp t_str $ BasicOp $ BinOp (SMin Int64) t d_K + insertTranspose :: + VarianceTable -> + (VName, SubExp) -> + (M.Map VName (Stm Kernels), M.Map VName (PrimType, Stm Kernels)) -> + (VName, Stm Kernels) -> + Binder Kernels (M.Map VName (Stm Kernels), M.Map VName (PrimType, Stm Kernels)) + insertTranspose variance (gidz, _) (tab_inn, tab_out) (p_nm, stm@(Let patt yy (BasicOp (Index arr_nm slc)))) + | [p] <- patternValueElements patt, + ptp <- elemType $ patElemType p, + p_nm == patElemName p = + case L.findIndices (variantSliceDim variance gidz) slc of + [] -> return (M.insert p_nm stm tab_inn, tab_out) + i : _ -> do + arr_tp <- lookupType arr_nm + let perm = [i + 1 .. arrayRank arr_tp -1] ++ [0 .. i] + let arr_tr_str = baseString arr_nm ++ "_transp" + arr_tr_nm <- letExp arr_tr_str $ BasicOp $ Manifest perm arr_nm + let e_ind' = BasicOp $ Index arr_tr_nm slc + let stm' = Let patt yy e_ind' + return (tab_inn, M.insert p_nm (ptp, stm') tab_out) + insertTranspose _ _ _ _ = error "\nUnreachable case reached in insertTranspose case, doRegTiling3D\n" + + variantSliceDim :: VarianceTable -> VName -> DimIndex SubExp -> Bool + variantSliceDim variance gidz (DimFix (Var vnm)) = variantToDim variance gidz vnm + variantSliceDim _ _ _ = False +doRegTiling3D _ = return Nothing 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 @@ -41,7 +41,7 @@ import Futhark.IR import Futhark.IR.Aliases ( Aliases, - consumedInStms, + mkStmsAliases, removeFunDefAliases, removeProgAliases, removeStmAliases, @@ -54,6 +54,9 @@ import Futhark.Pass import Futhark.Transform.Substitute +consumedInStms :: Aliased lore => Stms lore -> 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 @@ -136,7 +139,16 @@ runReader (cseInBody ds $ funDefBody fundec) $ newCSEState cse_arrays } where - ds = map (diet . declExtTypeOf) $ funDefRetType fundec + -- XXX: we treat every result as a consumption here, because we + -- our core language is not strong enough to fully capture the + -- aliases we want, so we are turning some parts off (see #803, + -- #1241, and the related comment in TypeCheck.hs). This is not a + -- practical problem while we still perform such aggressive + -- inlining. + ds = map retDiet $ funDefRetType fundec + retDiet t + | primType $ declExtTypeOf t = Observe + | otherwise = Consume type CSEM lore = Reader (CSEState lore) @@ -145,15 +157,16 @@ [Diet] -> Body lore -> CSEM lore (Body lore) -cseInBody ds (Body bodydec bnds res) = do - (bnds', res') <- - cseInStms (res_cons <> consumedInStms bnds) (stmsToList bnds) $ do +cseInBody ds (Body bodydec stms res) = do + (stms', res') <- + cseInStms (res_cons <> stms_cons) (stmsToList stms) $ do CSEState (_, nsubsts) _ <- ask return $ substituteNames nsubsts res - return $ Body bodydec bnds' res' + return $ Body bodydec stms' res' where - res_cons = mconcat $ zipWith consumeResult ds res - consumeResult Consume se = freeIn se + (res_als, stms_cons) = mkStmsAliases stms res + res_cons = mconcat $ zipWith consumeResult ds res_als + consumeResult Consume als = als consumeResult _ _ = mempty cseInLambda :: 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 @@ -528,7 +528,7 @@ -- the accumulator result cannot be fused! SOAC.Screma _ (ScremaForm scans reds _) _ -> drop (scanResults scans + redResults reds) out_nms - SOAC.Stream _ frm _ _ -> drop (length $ getStreamAccums frm) out_nms + SOAC.Stream _ _ _ nes _ -> drop (length nes) out_nms _ -> out_nms to_fuse_knms1 = S.toList $ getKersWithInpArrs res (out_arr_nms ++ inp_nms) to_fuse_knms2 = getKersWithSameInpSize (SOAC.width soac) res @@ -727,7 +727,7 @@ lambdaBody = lam_body, lambdaReturnType = map paramType $ acc_params ++ [offset_param] } - stream = Futhark.Stream w (Sequential $ merge_init ++ [intConst it 0]) lam loop_arrs + stream = Futhark.Stream w Sequential lam (merge_init ++ [intConst it 0]) loop_arrs -- It is important that the (discarded) final-offset is not the -- first element in the pattern, as we use the first element to @@ -757,16 +757,16 @@ Right soac@(SOAC.Screma _ (ScremaForm scans reds map_lam) _) -> reduceLike soac (map scanLambda scans <> map redLambda reds <> [map_lam]) $ concatMap scanNeutral scans <> concatMap redNeutral reds - Right soac@(SOAC.Stream _ form lam _) -> do + Right soac@(SOAC.Stream _ form lam nes _) -> do -- a redomap does not neccessarily start a new kernel, e.g., -- @let a= reduce(+,0,A) in ... bnds ... in let B = map(f,A)@ -- can be fused into a redomap that replaces the @map@, if @a@ -- and @B@ are defined in the same scope and @bnds@ does not uses @a@. -- a redomap always starts a new kernel let lambdas = case form of - Parallel _ _ lout _ -> [lout, lam] - _ -> [lam] - reduceLike soac lambdas $ getStreamAccums form + Parallel _ _ lout -> [lout, lam] + Sequential -> [lam] + reduceLike soac lambdas nes _ | [pe] <- patternValueElements pat, Just (src, trns) <- SOAC.transformFromExp (stmCerts bnd) e -> @@ -953,7 +953,7 @@ f_soac <- SOAC.toSOAC new_soac' -- The fused kernel may consume more than the original SOACs (see -- issue #224). We insert copy expressions to fix it. - f_soac' <- copyNewlyConsumed (fusedConsumed ker) $ addOpAliases f_soac + f_soac' <- copyNewlyConsumed (fusedConsumed ker) $ addOpAliases mempty f_soac validents <- zipWithM newIdent (map baseString names) $ SOAC.typeOf new_soac' auxing (kerAux ker <> aux) $ letBind (basicPattern [] validents) $ Op f_soac' transformOutput (outputTransform ker) names validents @@ -980,9 +980,9 @@ SOAC.Hist w ops lam arrs -> do lam' <- simplifyAndFuseInLambda lam return $ SOAC.Hist w ops lam' arrs - SOAC.Stream w form lam inps -> do + SOAC.Stream w form lam nes inps -> do lam' <- simplifyAndFuseInLambda lam - return $ SOAC.Stream w form lam' inps + return $ SOAC.Stream w form lam' nes inps simplifyAndFuseInLambda :: Lambda -> FusionGM Lambda simplifyAndFuseInLambda lam = do diff --git a/src/Futhark/Optimise/Fusion/LoopKernel.hs b/src/Futhark/Optimise/Fusion/LoopKernel.hs --- a/src/Futhark/Optimise/Fusion/LoopKernel.hs +++ b/src/Futhark/Optimise/Fusion/LoopKernel.hs @@ -484,19 +484,19 @@ ---------------------------- -- Stream-Stream Fusions: -- ---------------------------- - (SOAC.Stream _ Sequential {} _ _, SOAC.Stream _ form_p@Sequential {} _ _) - | mapFusionOK (drop (length $ getStreamAccums form_p) outVars) ker || horizFuse -> do + (SOAC.Stream _ Sequential _ _ _, SOAC.Stream _ Sequential _ nes _) + | mapFusionOK (drop (length nes) outVars) ker || horizFuse -> do -- fuse two SEQUENTIAL streams (res_nms, res_stream) <- fuseStreamHelper (outNames ker) unfus_set outVars outPairs soac_c soac_p success res_nms res_stream - (SOAC.Stream _ Sequential {} _ _, SOAC.Stream _ Sequential {} _ _) -> + (SOAC.Stream _ Sequential _ _ _, SOAC.Stream _ Sequential _ _ _) -> fail "Fusion conditions not met for two SEQ streams!" - (SOAC.Stream _ Sequential {} _ _, SOAC.Stream {}) -> + (SOAC.Stream _ Sequential _ _ _, SOAC.Stream {}) -> fail "Cannot fuse a parallel with a sequential Stream!" - (SOAC.Stream {}, SOAC.Stream _ Sequential {} _ _) -> + (SOAC.Stream {}, SOAC.Stream _ Sequential _ _ _) -> fail "Cannot fuse a parallel with a sequential Stream!" - (SOAC.Stream {}, SOAC.Stream _ form_p _ _) - | mapFusionOK (drop (length $ getStreamAccums form_p) outVars) ker || horizFuse -> do + (SOAC.Stream {}, SOAC.Stream _ _ _ nes _) + | mapFusionOK (drop (length nes) outVars) ker || horizFuse -> do -- fuse two PARALLEL streams (res_nms, res_stream) <- fuseStreamHelper (outNames ker) unfus_set outVars outPairs soac_c soac_p success res_nms res_stream @@ -511,7 +511,7 @@ --- we could run in an infinite recursion, i.e., repeatedly --- --- fusing map o scan into an infinity of Stream levels! --- ------------------------------------------------------------------- - (SOAC.Stream _ form2 _ _, _) -> do + (SOAC.Stream _ form2 _ _ _, _) -> do -- If this rule is matched then soac_p is NOT a stream. -- To fuse a stream kernel, we transform soac_p to a stream, which -- borrows the sequential/parallel property of the soac_c Stream, @@ -531,7 +531,7 @@ if soac_p' /= soac_p then fuseSOACwithKer unfus_set (map identName newacc_ids ++ outVars) soac_p' soac_p_consumed ker else fail "SOAC could not be turned into stream." - (_, SOAC.Stream _ form_p _ _) -> do + (_, SOAC.Stream _ form_p _ _ _) -> do -- If it reached this case then soac_c is NOT a Stream kernel, -- hence transform the kernel's soac to a stream and attempt -- stream-stream fusion recursivelly. @@ -540,7 +540,7 @@ (soac_c', newacc_ids) <- SOAC.soacToStream soac_c when (soac_c' == soac_c) $ fail "SOAC could not be turned into stream." soac_c'' <- case form_p of - Sequential _ -> toSeqStream soac_c' + Sequential -> toSeqStream soac_c' _ -> return soac_c' fuseSOACwithKer unfus_set outVars soac_p soac_p_consumed $ @@ -552,8 +552,8 @@ _ -> fail "Cannot fuse" getStreamOrder :: StreamForm lore -> StreamOrd -getStreamOrder (Parallel o _ _ _) = o -getStreamOrder (Sequential _) = InOrder +getStreamOrder (Parallel o _ _) = o +getStreamOrder Sequential = InOrder fuseStreamHelper :: [VName] -> @@ -568,16 +568,15 @@ unfus_set outVars outPairs - (SOAC.Stream w2 form2 lam2 inp2_arr) - (SOAC.Stream _ form1 lam1 inp1_arr) = + (SOAC.Stream w2 form2 lam2 nes2 inp2_arr) + (SOAC.Stream _ form1 lam1 nes1 inp1_arr) = if getStreamOrder form2 /= getStreamOrder form1 then fail "fusion conditions not met!" else do -- very similar to redomap o redomap composition, but need -- to remove first the `chunk' parameters of streams' -- lambdas and put them in the resulting stream lambda. - let nes1 = getStreamAccums form1 - chunk1 = head $ lambdaParams lam1 + let chunk1 = head $ lambdaParams lam1 chunk2 = head $ lambdaParams lam2 hmnms = M.fromList [(paramName chunk2, paramName chunk1)] lam20 = substituteNames hmnms lam2 @@ -594,7 +593,7 @@ outPairs lam2' [] - (getStreamAccums form2) + nes2 inp2_arr res_lam'' = res_lam' {lambdaParams = chunk1 : lambdaParams res_lam'} unfus_accs = take (length nes1) outVars @@ -602,21 +601,21 @@ res_form <- mergeForms form2 form1 return ( unfus_accs ++ out_kernms ++ unfus_arrs, - SOAC.Stream w2 res_form res_lam'' new_inp + SOAC.Stream w2 res_form res_lam'' (nes1 ++ nes2) new_inp ) where - mergeForms (Sequential acc2) (Sequential acc1) = return $ Sequential (acc1 ++ acc2) - mergeForms (Parallel _ comm2 lam2r acc2) (Parallel o1 comm1 lam1r acc1) = - return $ Parallel o1 (comm1 <> comm2) (mergeReduceOps lam1r lam2r) (acc1 ++ acc2) + mergeForms Sequential Sequential = return Sequential + mergeForms (Parallel _ comm2 lam2r) (Parallel o1 comm1 lam1r) = + return $ Parallel o1 (comm1 <> comm2) (mergeReduceOps lam1r lam2r) mergeForms _ _ = fail "Fusing sequential to parallel stream disallowed!" fuseStreamHelper _ _ _ _ _ _ = fail "Cannot Fuse Streams!" -- | If a Stream is passed as argument then it converts it to a -- Sequential Stream; Otherwise it FAILS! toSeqStream :: SOAC -> TryFusion SOAC -toSeqStream s@(SOAC.Stream _ (Sequential _) _ _) = return s -toSeqStream (SOAC.Stream w (Parallel _ _ _ acc) l inps) = - return $ SOAC.Stream w (Sequential acc) l inps +toSeqStream s@(SOAC.Stream _ Sequential _ _ _) = return s +toSeqStream (SOAC.Stream w Parallel {} l acc inps) = + return $ SOAC.Stream w Sequential l acc inps toSeqStream _ = fail "toSeqStream expects a stream, but given a SOAC." -- Here follows optimizations and transforms to expose fusability. diff --git a/src/Futhark/Optimise/Simplify/ClosedForm.hs b/src/Futhark/Optimise/Simplify/ClosedForm.hs deleted file mode 100644 --- a/src/Futhark/Optimise/Simplify/ClosedForm.hs +++ /dev/null @@ -1,226 +0,0 @@ -{-# LANGUAGE FlexibleContexts #-} - --- | This module implements facilities for determining whether a --- reduction or fold can be expressed in a closed form (i.e. not as a --- SOAC). --- --- Right now, the module can detect only trivial cases. In the --- future, we would like to make it more powerful, as well as possibly --- also being able to analyse sequential loops. -module Futhark.Optimise.Simplify.ClosedForm - ( foldClosedForm, - loopClosedForm, - VarLookup, - ) -where - -import Control.Monad -import qualified Data.Map.Strict as M -import Data.Maybe -import Futhark.Construct -import Futhark.IR -import Futhark.Optimise.Simplify.Rule -import Futhark.Transform.Rename - --- | A function that, given a variable name, returns its definition. -type VarLookup lore = VName -> Maybe (Exp lore, Certificates) - -{- -Motivation: - - let {*[int,x_size_27] map_computed_shape_1286} = replicate(x_size_27, - all_equal_shape_1044) in - let {*[bool,x_size_27] map_size_checks_1292} = replicate(x_size_27, x_1291) in - let {bool all_equal_checked_1298, int all_equal_shape_1299} = - reduceT(fn {bool, int} (bool bacc_1293, int nacc_1294, bool belm_1295, - int nelm_1296) => - let {bool tuplit_elems_1297} = bacc_1293 && belm_1295 in - {tuplit_elems_1297, nelm_1296}, - {True, 0}, map_size_checks_1292, map_computed_shape_1286) --} - --- | @foldClosedForm look foldfun accargs arrargs@ determines whether --- each of the results of @foldfun@ can be expressed in a closed form. -foldClosedForm :: - (ASTLore lore, BinderOps lore) => - VarLookup lore -> - Pattern lore -> - Lambda lore -> - [SubExp] -> - [VName] -> - RuleM lore () -foldClosedForm look pat lam accs arrs = do - inputsize <- arraysSize 0 <$> mapM lookupType arrs - - t <- case patternTypes pat of - [Prim t] -> return t - _ -> cannotSimplify - - closedBody <- - checkResults - (patternNames pat) - inputsize - mempty - Int64 - knownBnds - (map paramName (lambdaParams lam)) - (lambdaBody lam) - accs - isEmpty <- newVName "fold_input_is_empty" - letBindNames [isEmpty] $ - BasicOp $ CmpOp (CmpEq int64) inputsize (intConst Int64 0) - letBind pat - =<< ( If (Var isEmpty) - <$> resultBodyM accs - <*> renameBody closedBody - <*> pure (IfDec [primBodyType t] IfNormal) - ) - where - knownBnds = determineKnownBindings look lam accs arrs - --- | @loopClosedForm pat respat merge bound bodys@ determines whether --- the do-loop can be expressed in a closed form. -loopClosedForm :: - (ASTLore lore, BinderOps lore) => - Pattern lore -> - [(FParam lore, SubExp)] -> - Names -> - IntType -> - SubExp -> - Body lore -> - RuleM lore () -loopClosedForm pat merge i it bound body = do - t <- case patternTypes pat of - [Prim t] -> return t - _ -> cannotSimplify - - closedBody <- - checkResults - mergenames - bound - i - it - knownBnds - (map identName mergeidents) - body - mergeexp - isEmpty <- newVName "bound_is_zero" - letBindNames [isEmpty] $ - BasicOp $ CmpOp (CmpSlt it) bound (intConst it 0) - - letBind pat - =<< ( If (Var isEmpty) - <$> resultBodyM mergeexp - <*> renameBody closedBody - <*> pure (IfDec [primBodyType t] IfNormal) - ) - where - (mergepat, mergeexp) = unzip merge - mergeidents = map paramIdent mergepat - mergenames = map paramName mergepat - knownBnds = M.fromList $ zip mergenames mergeexp - -checkResults :: - BinderOps lore => - [VName] -> - SubExp -> - Names -> - IntType -> - M.Map VName SubExp -> - -- | Lambda-bound - [VName] -> - Body lore -> - [SubExp] -> - RuleM lore (Body lore) -checkResults pat size untouchable it knownBnds params body accs = do - ((), bnds) <- - collectStms $ - zipWithM_ checkResult (zip pat res) (zip accparams accs) - mkBodyM bnds $ map Var pat - where - bndMap = makeBindMap body - (accparams, _) = splitAt (length accs) params - res = bodyResult body - - nonFree = boundInBody body <> namesFromList params <> untouchable - - checkResult (p, Var v) (accparam, acc) - | Just (BasicOp (BinOp bop x y)) <- M.lookup v bndMap = do - -- One of x,y must be *this* accumulator, and the other must - -- be something that is free in the body. - let isThisAccum = (== Var accparam) - (this, el) <- liftMaybe $ - case ( (asFreeSubExp x, isThisAccum y), - (asFreeSubExp y, isThisAccum x) - ) of - ((Just free, True), _) -> Just (acc, free) - (_, (Just free, True)) -> Just (acc, free) - _ -> Nothing - - case bop of - LogAnd -> - letBindNames [p] $ BasicOp $ BinOp LogAnd this el - Add t w | Just properly_typed_size <- properIntSize t -> do - size' <- properly_typed_size - letBindNames [p] - =<< eBinOp - (Add t w) - (eSubExp this) - (pure $ BasicOp $ BinOp (Mul t w) el size') - FAdd t | Just properly_typed_size <- properFloatSize t -> do - size' <- properly_typed_size - letBindNames [p] - =<< eBinOp - (FAdd t) - (eSubExp this) - (pure $ BasicOp $ BinOp (FMul t) el size') - _ -> cannotSimplify -- Um... sorry. - checkResult _ _ = cannotSimplify - - asFreeSubExp :: SubExp -> Maybe SubExp - asFreeSubExp (Var v) - | v `nameIn` nonFree = M.lookup v knownBnds - asFreeSubExp se = Just se - - properIntSize Int64 = Just $ return size - properIntSize t = - Just $ - letSubExp "converted_size" $ - BasicOp $ ConvOp (SExt it t) size - - properFloatSize t = - Just $ - letSubExp "converted_size" $ - BasicOp $ ConvOp (SIToFP it t) size - -determineKnownBindings :: - VarLookup lore -> - Lambda lore -> - [SubExp] -> - [VName] -> - M.Map VName SubExp -determineKnownBindings look lam accs arrs = - accBnds <> arrBnds - where - (accparams, arrparams) = - splitAt (length accs) $ lambdaParams lam - accBnds = - M.fromList $ - zip (map paramName accparams) accs - arrBnds = - M.fromList $ - mapMaybe isReplicate $ - zip (map paramName arrparams) arrs - - isReplicate (p, v) - | Just (BasicOp (Replicate _ ve), cs) <- look v, - cs == mempty = - Just (p, ve) - isReplicate _ = Nothing - -makeBindMap :: Body lore -> M.Map VName (Exp lore) -makeBindMap = M.fromList . mapMaybe isSingletonStm . stmsToList . bodyStms - where - isSingletonStm (Let pat _ e) = case patternNames pat of - [v] -> Just (v, e) - _ -> Nothing diff --git a/src/Futhark/Optimise/Simplify/Rules.hs b/src/Futhark/Optimise/Simplify/Rules.hs --- a/src/Futhark/Optimise/Simplify/Rules.hs +++ b/src/Futhark/Optimise/Simplify/Rules.hs @@ -1,1468 +1,268 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RankNTypes #-} -{-# LANGUAGE TupleSections #-} -{-# LANGUAGE TypeFamilies #-} - --- | This module defines a collection of simplification rules, as per --- "Futhark.Optimise.Simplify.Rule". They are used in the --- simplifier. --- --- For performance reasons, many sufficiently simple logically --- separate rules are merged into single "super-rules", like ruleIf --- and ruleBasicOp. This is because it is relatively expensive to --- activate a rule just to determine that it does not apply. Thus, it --- is more efficient to have a few very fat rules than a lot of small --- rules. This does not affect the compiler result in any way; it is --- purely an optimisation to speed up compilation. -module Futhark.Optimise.Simplify.Rules - ( standardRules, - removeUnnecessaryCopy, - ) -where - -import Control.Monad -import Data.Either -import Data.List (find, foldl', isSuffixOf, partition, sort) -import qualified Data.Map.Strict as M -import Data.Maybe -import Futhark.Analysis.DataDependencies -import Futhark.Analysis.PrimExp.Convert -import qualified Futhark.Analysis.SymbolTable as ST -import qualified Futhark.Analysis.UsageTable as UT -import Futhark.Construct -import Futhark.IR -import Futhark.IR.Prop.Aliases -import Futhark.Optimise.Simplify.ClosedForm -import Futhark.Optimise.Simplify.Rule -import Futhark.Transform.Rename -import Futhark.Util - -topDownRules :: (BinderOps lore, Aliased lore) => [TopDownRule lore] -topDownRules = - [ RuleDoLoop hoistLoopInvariantMergeVariables, - RuleDoLoop simplifyClosedFormLoop, - RuleDoLoop simplifyKnownIterationLoop, - RuleDoLoop simplifyLoopVariables, - RuleDoLoop narrowLoopType, - RuleGeneric constantFoldPrimFun, - RuleIf ruleIf, - RuleIf hoistBranchInvariant, - RuleBasicOp ruleBasicOp - ] - -bottomUpRules :: BinderOps lore => [BottomUpRule lore] -bottomUpRules = - [ RuleDoLoop removeRedundantMergeVariables, - RuleIf removeDeadBranchResult, - RuleBasicOp simplifyIndex, - RuleBasicOp simplifyConcat - ] - --- | A set of standard simplification rules. These assume pure --- functional semantics, and so probably should not be applied after --- memory block merging. -standardRules :: (BinderOps lore, Aliased lore) => RuleBook lore -standardRules = ruleBook topDownRules bottomUpRules - --- This next one is tricky - it's easy enough to determine that some --- loop result is not used after the loop, but here, we must also make --- sure that it does not affect any other values. --- --- I do not claim that the current implementation of this rule is --- perfect, but it should suffice for many cases, and should never --- generate wrong code. -removeRedundantMergeVariables :: BinderOps lore => BottomUpRuleDoLoop lore -removeRedundantMergeVariables (_, used) pat aux (ctx, val, form, body) - | not $ all (usedAfterLoop . fst) val, - null ctx -- FIXME: things get tricky if we can remove all vals - -- but some ctxs are still used. We take the easy way - -- out for now. - = - let (ctx_es, val_es) = splitAt (length ctx) $ bodyResult body - necessaryForReturned = - findNecessaryForReturned - usedAfterLoopOrInForm - (zip (map fst $ ctx ++ val) $ ctx_es ++ val_es) - (dataDependencies body) - - resIsNecessary ((v, _), _) = - usedAfterLoop v - || paramName v `nameIn` necessaryForReturned - || referencedInPat v - || referencedInForm v - - (keep_ctx, discard_ctx) = - partition resIsNecessary $ zip ctx ctx_es - (keep_valpart, discard_valpart) = - partition (resIsNecessary . snd) $ - zip (patternValueElements pat) $ zip val val_es - - (keep_valpatelems, keep_val) = unzip keep_valpart - (_discard_valpatelems, discard_val) = unzip discard_valpart - (ctx', ctx_es') = unzip keep_ctx - (val', val_es') = unzip keep_val - - body' = body {bodyResult = ctx_es' ++ val_es'} - free_in_keeps = freeIn keep_valpatelems - - stillUsedContext pat_elem = - patElemName pat_elem - `nameIn` ( free_in_keeps - <> freeIn (filter (/= pat_elem) $ patternContextElements pat) - ) - - pat' = - pat - { patternValueElements = keep_valpatelems, - patternContextElements = - filter stillUsedContext $ patternContextElements pat - } - in if ctx' ++ val' == ctx ++ val - then Skip - else Simplify $ do - -- We can't just remove the bindings in 'discard', since the loop - -- body may still use their names in (now-dead) expressions. - -- Hence, we add them inside the loop, fully aware that dead-code - -- removal will eventually get rid of them. Some care is - -- necessary to handle unique bindings. - body'' <- insertStmsM $ do - mapM_ (uncurry letBindNames) $ dummyStms discard_ctx - mapM_ (uncurry letBindNames) $ dummyStms discard_val - return body' - auxing aux $ letBind pat' $ DoLoop ctx' val' form body'' - where - pat_used = map (`UT.isUsedDirectly` used) $ patternValueNames pat - used_vals = map fst $ filter snd $ zip (map (paramName . fst) val) pat_used - usedAfterLoop = flip elem used_vals . paramName - usedAfterLoopOrInForm p = - usedAfterLoop p || paramName p `nameIn` freeIn form - patAnnotNames = freeIn $ map fst $ ctx ++ val - referencedInPat = (`nameIn` patAnnotNames) . paramName - referencedInForm = (`nameIn` freeIn form) . paramName - - dummyStms = map dummyStm - dummyStm ((p, e), _) - | unique (paramDeclType p), - Var v <- e = - ([paramName p], BasicOp $ Copy v) - | otherwise = ([paramName p], BasicOp $ SubExp e) -removeRedundantMergeVariables _ _ _ _ = - Skip - --- We may change the type of the loop if we hoist out a shape --- annotation, in which case we also need to tweak the bound pattern. -hoistLoopInvariantMergeVariables :: BinderOps lore => TopDownRuleDoLoop lore -hoistLoopInvariantMergeVariables vtable pat aux (ctx, val, form, loopbody) = - -- Figure out which of the elements of loopresult are - -- loop-invariant, and hoist them out. - case foldr checkInvariance ([], explpat, [], []) $ - zip3 (patternNames pat) merge res of - ([], _, _, _) -> - -- Nothing is invariant. - Skip - (invariant, explpat', merge', res') -> Simplify $ do - -- We have moved something invariant out of the loop. - let loopbody' = loopbody {bodyResult = res'} - invariantShape :: (a, VName) -> Bool - invariantShape (_, shapemerge) = - shapemerge - `elem` map (paramName . fst) merge' - (implpat', implinvariant) = partition invariantShape implpat - implinvariant' = [(patElemIdent p, Var v) | (p, v) <- implinvariant] - implpat'' = map fst implpat' - explpat'' = map fst explpat' - (ctx', val') = splitAt (length implpat') merge' - forM_ (invariant ++ implinvariant') $ \(v1, v2) -> - letBindNames [identName v1] $ BasicOp $ SubExp v2 - auxing aux $ - letBind (Pattern implpat'' explpat'') $ - DoLoop ctx' val' form loopbody' - where - merge = ctx ++ val - res = bodyResult loopbody - - implpat = - zip (patternContextElements pat) $ - map (paramName . fst) ctx - explpat = - zip (patternValueElements pat) $ - map (paramName . fst) val - - namesOfMergeParams = namesFromList $ map (paramName . fst) $ ctx ++ val - - removeFromResult (mergeParam, mergeInit) explpat' = - case partition ((== paramName mergeParam) . snd) explpat' of - ([(patelem, _)], rest) -> - (Just (patElemIdent patelem, mergeInit), rest) - (_, _) -> - (Nothing, explpat') - - checkInvariance - (pat_name, (mergeParam, mergeInit), resExp) - (invariant, explpat', merge', resExps) - | not (unique (paramDeclType mergeParam)) - || arrayRank (paramDeclType mergeParam) == 1, - isInvariant, - -- Also do not remove the condition in a while-loop. - not $ paramName mergeParam `nameIn` freeIn form = - let (bnd, explpat'') = - removeFromResult (mergeParam, mergeInit) explpat' - in ( maybe id (:) bnd $ (paramIdent mergeParam, mergeInit) : invariant, - explpat'', - merge', - resExps - ) - where - -- A non-unique merge variable is invariant if one of the - -- following is true: - -- - -- (0) The result is a variable of the same name as the - -- parameter, where all existential parameters are already - -- known to be invariant - isInvariant - | Var v2 <- resExp, - paramName mergeParam == v2 = - allExistentialInvariant - (namesFromList $ map (identName . fst) invariant) - mergeParam - -- (1) The result is identical to the initial parameter value. - | mergeInit == resExp = True - -- (2) The initial parameter value is equal to an outer - -- loop parameter 'P', where the initial value of 'P' is - -- equal to 'resExp', AND 'resExp' ultimately becomes the - -- new value of 'P'. XXX: it's a bit clumsy that this - -- only works for one level of nesting, and I think it - -- would not be too hard to generalise. - | Var init_v <- mergeInit, - Just (p_init, p_res) <- ST.lookupLoopParam init_v vtable, - p_init == resExp, - p_res == Var pat_name = - True - | otherwise = False - checkInvariance - (_pat_name, (mergeParam, mergeInit), resExp) - (invariant, explpat', merge', resExps) = - (invariant, explpat', (mergeParam, mergeInit) : merge', resExp : resExps) - - allExistentialInvariant namesOfInvariant mergeParam = - all (invariantOrNotMergeParam namesOfInvariant) $ - namesToList $ - freeIn mergeParam `namesSubtract` oneName (paramName mergeParam) - invariantOrNotMergeParam namesOfInvariant name = - not (name `nameIn` namesOfMergeParams) - || name `nameIn` namesOfInvariant - --- | A function that, given a subexpression, returns its type. -type TypeLookup = SubExp -> Maybe Type - --- | A simple rule is a top-down rule that can be expressed as a pure --- function. -type SimpleRule lore = VarLookup lore -> TypeLookup -> BasicOp -> Maybe (BasicOp, Certificates) - -simpleRules :: [SimpleRule lore] -simpleRules = - [ simplifyBinOp, - simplifyCmpOp, - simplifyUnOp, - simplifyConvOp, - simplifyAssert, - copyScratchToScratch, - simplifyIdentityReshape, - simplifyReshapeReshape, - simplifyReshapeScratch, - simplifyReshapeReplicate, - simplifyReshapeIota, - improveReshape - ] - -simplifyClosedFormLoop :: BinderOps lore => TopDownRuleDoLoop lore -simplifyClosedFormLoop _ pat _ ([], val, ForLoop i it bound [], body) = - Simplify $ loopClosedForm pat val (oneName i) it bound body -simplifyClosedFormLoop _ _ _ _ = Skip - -simplifyLoopVariables :: (BinderOps lore, Aliased lore) => TopDownRuleDoLoop lore -simplifyLoopVariables vtable pat aux (ctx, val, form@(ForLoop i it num_iters loop_vars), body) - | simplifiable <- map checkIfSimplifiable loop_vars, - not $ all isNothing simplifiable = Simplify $ do - -- Check if the simplifications throw away more information than - -- we are comfortable with at this stage. - (maybe_loop_vars, body_prefix_stms) <- - localScope (scopeOf form) $ - unzip <$> zipWithM onLoopVar loop_vars simplifiable - if maybe_loop_vars == map Just loop_vars - then cannotSimplify - else do - body' <- insertStmsM $ do - addStms $ mconcat body_prefix_stms - resultBodyM =<< bodyBind body - auxing aux $ - letBind pat $ - DoLoop - ctx - val - (ForLoop i it num_iters $ catMaybes maybe_loop_vars) - body' - where - seType (Var v) - | v == i = Just $ Prim $ IntType it - | otherwise = ST.lookupType v vtable - seType (Constant v) = Just $ Prim $ primValueType v - consumed_in_body = consumedInBody body - - vtable' = ST.fromScope (scopeOf form) <> vtable - - checkIfSimplifiable (p, arr) = - simplifyIndexing - vtable' - seType - arr - (DimFix (Var i) : fullSlice (paramType p) []) - $ paramName p `nameIn` consumed_in_body - - -- We only want this simplification if the result does not refer - -- to 'i' at all, or does not contain accesses. - onLoopVar (p, arr) Nothing = - return (Just (p, arr), mempty) - onLoopVar (p, arr) (Just m) = do - (x, x_stms) <- collectStms m - case x of - IndexResult cs arr' slice - | not $ any ((i `nameIn`) . freeIn) x_stms, - DimFix (Var j) : slice' <- slice, - j == i, - not $ i `nameIn` freeIn slice -> do - addStms x_stms - w <- arraySize 0 <$> lookupType arr' - for_in_partial <- - certifying cs $ - letExp "for_in_partial" $ - BasicOp $ - Index arr' $ - DimSlice (intConst Int64 0) w (intConst Int64 1) : slice' - return (Just (p, for_in_partial), mempty) - SubExpResult cs se - | all (notIndex . stmExp) x_stms -> do - x_stms' <- collectStms_ $ - certifying cs $ do - addStms x_stms - letBindNames [paramName p] $ BasicOp $ SubExp se - return (Nothing, x_stms') - _ -> return (Just (p, arr), mempty) - - notIndex (BasicOp Index {}) = False - notIndex _ = True -simplifyLoopVariables _ _ _ _ = Skip - --- If a for-loop with no loop variables has a counter of type Int64, --- and the bound is just a constant or sign-extended integer of --- smaller type, then change the loop to iterate over the smaller type --- instead. We then move the sign extension inside the loop instead. --- This addresses loops of the form @for i in x.. TopDownRuleDoLoop lore -narrowLoopType vtable pat aux (ctx, val, ForLoop i Int64 n [], body) - | Just (n', it', cs) <- smallerType = - Simplify $ do - i' <- newVName $ baseString i - let form' = ForLoop i' it' n' [] - body' <- insertStmsM $ - inScopeOf form' $ do - letBindNames [i] $ BasicOp $ ConvOp (SExt it' Int64) (Var i') - pure body - auxing aux $ - certifying cs $ - letBind pat $ DoLoop ctx val form' body' - where - smallerType - | Var n' <- n, - Just (ConvOp (SExt it' _) n'', cs) <- ST.lookupBasicOp n' vtable = - Just (n'', it', cs) - | Constant (IntValue (Int64Value n')) <- n, - toInteger n' <= toInteger (maxBound :: Int32) = - Just (intConst Int32 (toInteger n'), Int32, mempty) - | otherwise = - Nothing -narrowLoopType _ _ _ _ = Skip - -unroll :: - BinderOps lore => - Integer -> - [(FParam lore, SubExp)] -> - (VName, IntType, Integer) -> - [(LParam lore, VName)] -> - Body lore -> - RuleM lore [SubExp] -unroll n merge (iv, it, i) loop_vars body - | i >= n = - return $ map snd merge - | otherwise = do - iter_body <- insertStmsM $ do - forM_ merge $ \(mergevar, mergeinit) -> - letBindNames [paramName mergevar] $ BasicOp $ SubExp mergeinit - - letBindNames [iv] $ BasicOp $ SubExp $ intConst it i - - forM_ loop_vars $ \(p, arr) -> - letBindNames [paramName p] $ - BasicOp $ - Index arr $ - DimFix (intConst Int64 i) : fullSlice (paramType p) [] - - -- Some of the sizes in the types here might be temporarily wrong - -- until copy propagation fixes it up. - pure body - - iter_body' <- renameBody iter_body - addStms $ bodyStms iter_body' - - let merge' = zip (map fst merge) $ bodyResult iter_body' - unroll n merge' (iv, it, i + 1) loop_vars body - -simplifyKnownIterationLoop :: BinderOps lore => TopDownRuleDoLoop lore -simplifyKnownIterationLoop _ pat aux (ctx, val, ForLoop i it (Constant iters) loop_vars, body) - | IntValue n <- iters, - zeroIshInt n || oneIshInt n || "unroll" `inAttrs` stmAuxAttrs aux = Simplify $ do - res <- unroll (valueIntegral n) (ctx ++ val) (i, it, 0) loop_vars body - forM_ (zip (patternNames pat) res) $ \(v, se) -> - letBindNames [v] $ BasicOp $ SubExp se -simplifyKnownIterationLoop _ _ _ _ = - Skip - --- | Turn @copy(x)@ into @x@ iff @x@ is not used after this copy --- statement and it can be consumed. --- --- This simplistic rule is only valid before we introduce memory. -removeUnnecessaryCopy :: BinderOps lore => BottomUpRuleBasicOp lore -removeUnnecessaryCopy (vtable, used) (Pattern [] [d]) _ (Copy v) - | not (v `UT.isConsumed` used), - (not (v `UT.used` used) && consumable) || not (patElemName d `UT.isConsumed` used) = - Simplify $ letBindNames [patElemName d] $ BasicOp $ SubExp $ Var v - where - -- We need to make sure we can even consume the original. - -- This is currently a hacky check, much too conservative, - -- because we don't have the information conveniently - -- available. - consumable = case M.lookup v $ ST.toScope vtable of - Just (FParamName info) -> unique $ declTypeOf info - _ -> False -removeUnnecessaryCopy _ _ _ _ = Skip - -simplifyCmpOp :: SimpleRule lore -simplifyCmpOp _ _ (CmpOp cmp e1 e2) - | e1 == e2 = constRes $ - BoolValue $ - case cmp of - CmpEq {} -> True - CmpSlt {} -> False - CmpUlt {} -> False - CmpSle {} -> True - CmpUle {} -> True - FCmpLt {} -> False - FCmpLe {} -> True - CmpLlt -> False - CmpLle -> True -simplifyCmpOp _ _ (CmpOp cmp (Constant v1) (Constant v2)) = - constRes . BoolValue =<< doCmpOp cmp v1 v2 -simplifyCmpOp look _ (CmpOp CmpEq {} (Constant (IntValue x)) (Var v)) - | Just (BasicOp (ConvOp BToI {} b), cs) <- look v = - case valueIntegral x :: Int of - 1 -> Just (SubExp b, cs) - 0 -> Just (UnOp Not b, cs) - _ -> Just (SubExp (Constant (BoolValue False)), cs) -simplifyCmpOp _ _ _ = Nothing - -simplifyBinOp :: SimpleRule lore -simplifyBinOp _ _ (BinOp op (Constant v1) (Constant v2)) - | Just res <- doBinOp op v1 v2 = - constRes res -simplifyBinOp look _ (BinOp Add {} e1 e2) - | isCt0 e1 = subExpRes e2 - | isCt0 e2 = subExpRes e1 - -- x+(y-x) => y - | Var v2 <- e2, - Just (BasicOp (BinOp Sub {} e2_a e2_b), cs) <- look v2, - e2_b == e1 = - Just (SubExp e2_a, cs) -simplifyBinOp _ _ (BinOp FAdd {} e1 e2) - | isCt0 e1 = subExpRes e2 - | isCt0 e2 = subExpRes e1 -simplifyBinOp look _ (BinOp sub@(Sub t _) e1 e2) - | isCt0 e2 = subExpRes e1 - -- Cases for simplifying (a+b)-b and permutations. - - -- (e1_a+e1_b)-e1_a == e1_b - | Var v1 <- e1, - Just (BasicOp (BinOp Add {} e1_a e1_b), cs) <- look v1, - e1_a == e2 = - Just (SubExp e1_b, cs) - -- (e1_a+e1_b)-e1_b == e1_a - | Var v1 <- e1, - Just (BasicOp (BinOp Add {} e1_a e1_b), cs) <- look v1, - e1_b == e2 = - Just (SubExp e1_a, cs) - -- e2_a-(e2_a+e2_b) == 0-e2_b - | Var v2 <- e2, - Just (BasicOp (BinOp Add {} e2_a e2_b), cs) <- look v2, - e2_a == e1 = - Just (BinOp sub (intConst t 0) e2_b, cs) - -- e2_b-(e2_a+e2_b) == 0-e2_a - | Var v2 <- e2, - Just (BasicOp (BinOp Add {} e2_a e2_b), cs) <- look v2, - e2_b == e1 = - Just (BinOp sub (intConst t 0) e2_a, cs) -simplifyBinOp _ _ (BinOp FSub {} e1 e2) - | isCt0 e2 = subExpRes e1 -simplifyBinOp _ _ (BinOp Mul {} e1 e2) - | isCt0 e1 = subExpRes e1 - | isCt0 e2 = subExpRes e2 - | isCt1 e1 = subExpRes e2 - | isCt1 e2 = subExpRes e1 -simplifyBinOp _ _ (BinOp FMul {} e1 e2) - | isCt0 e1 = subExpRes e1 - | isCt0 e2 = subExpRes e2 - | isCt1 e1 = subExpRes e2 - | isCt1 e2 = subExpRes e1 -simplifyBinOp look _ (BinOp (SMod t _) e1 e2) - | isCt1 e2 = constRes $ IntValue $ intValue t (0 :: Int) - | e1 == e2 = constRes $ IntValue $ intValue t (0 :: Int) - | Var v1 <- e1, - Just (BasicOp (BinOp SMod {} _ e4), v1_cs) <- look v1, - e4 == e2 = - Just (SubExp e1, v1_cs) -simplifyBinOp _ _ (BinOp SDiv {} e1 e2) - | isCt0 e1 = subExpRes e1 - | isCt1 e2 = subExpRes e1 - | isCt0 e2 = Nothing -simplifyBinOp _ _ (BinOp SDivUp {} e1 e2) - | isCt0 e1 = subExpRes e1 - | isCt1 e2 = subExpRes e1 - | isCt0 e2 = Nothing -simplifyBinOp _ _ (BinOp FDiv {} e1 e2) - | isCt0 e1 = subExpRes e1 - | isCt1 e2 = subExpRes e1 - | isCt0 e2 = Nothing -simplifyBinOp _ _ (BinOp (SRem t _) e1 e2) - | isCt1 e2 = constRes $ IntValue $ intValue t (0 :: Int) - | e1 == e2 = constRes $ IntValue $ intValue t (1 :: Int) -simplifyBinOp _ _ (BinOp SQuot {} e1 e2) - | isCt1 e2 = subExpRes e1 - | isCt0 e2 = Nothing -simplifyBinOp _ _ (BinOp (FPow t) e1 e2) - | isCt0 e2 = subExpRes $ floatConst t 1 - | isCt0 e1 || isCt1 e1 || isCt1 e2 = subExpRes e1 -simplifyBinOp _ _ (BinOp (Shl t) e1 e2) - | isCt0 e2 = subExpRes e1 - | isCt0 e1 = subExpRes $ intConst t 0 -simplifyBinOp _ _ (BinOp AShr {} e1 e2) - | isCt0 e2 = subExpRes e1 -simplifyBinOp _ _ (BinOp (And t) e1 e2) - | isCt0 e1 = subExpRes $ intConst t 0 - | isCt0 e2 = subExpRes $ intConst t 0 - | e1 == e2 = subExpRes e1 -simplifyBinOp _ _ (BinOp Or {} e1 e2) - | isCt0 e1 = subExpRes e2 - | isCt0 e2 = subExpRes e1 - | e1 == e2 = subExpRes e1 -simplifyBinOp _ _ (BinOp (Xor t) e1 e2) - | isCt0 e1 = subExpRes e2 - | isCt0 e2 = subExpRes e1 - | e1 == e2 = subExpRes $ intConst t 0 -simplifyBinOp defOf _ (BinOp LogAnd e1 e2) - | isCt0 e1 = constRes $ BoolValue False - | isCt0 e2 = constRes $ BoolValue False - | isCt1 e1 = subExpRes e2 - | isCt1 e2 = subExpRes e1 - | Var v <- e1, - Just (BasicOp (UnOp Not e1'), v_cs) <- defOf v, - e1' == e2 = - Just (SubExp $ Constant $ BoolValue False, v_cs) - | Var v <- e2, - Just (BasicOp (UnOp Not e2'), v_cs) <- defOf v, - e2' == e1 = - Just (SubExp $ Constant $ BoolValue False, v_cs) -simplifyBinOp defOf _ (BinOp LogOr e1 e2) - | isCt0 e1 = subExpRes e2 - | isCt0 e2 = subExpRes e1 - | isCt1 e1 = constRes $ BoolValue True - | isCt1 e2 = constRes $ BoolValue True - | Var v <- e1, - Just (BasicOp (UnOp Not e1'), v_cs) <- defOf v, - e1' == e2 = - Just (SubExp $ Constant $ BoolValue True, v_cs) - | Var v <- e2, - Just (BasicOp (UnOp Not e2'), v_cs) <- defOf v, - e2' == e1 = - Just (SubExp $ Constant $ BoolValue True, v_cs) -simplifyBinOp defOf _ (BinOp (SMax it) e1 e2) - | e1 == e2 = - subExpRes e1 - | Var v1 <- e1, - Just (BasicOp (BinOp (SMax _) e1_1 e1_2), v1_cs) <- defOf v1, - e1_1 == e2 = - Just (BinOp (SMax it) e1_2 e2, v1_cs) - | Var v1 <- e1, - Just (BasicOp (BinOp (SMax _) e1_1 e1_2), v1_cs) <- defOf v1, - e1_2 == e2 = - Just (BinOp (SMax it) e1_1 e2, v1_cs) - | Var v2 <- e2, - Just (BasicOp (BinOp (SMax _) e2_1 e2_2), v2_cs) <- defOf v2, - e2_1 == e1 = - Just (BinOp (SMax it) e2_2 e1, v2_cs) - | Var v2 <- e2, - Just (BasicOp (BinOp (SMax _) e2_1 e2_2), v2_cs) <- defOf v2, - e2_2 == e1 = - Just (BinOp (SMax it) e2_1 e1, v2_cs) -simplifyBinOp _ _ _ = Nothing - -constRes :: PrimValue -> Maybe (BasicOp, Certificates) -constRes = Just . (,mempty) . SubExp . Constant - -subExpRes :: SubExp -> Maybe (BasicOp, Certificates) -subExpRes = Just . (,mempty) . SubExp - -simplifyUnOp :: SimpleRule lore -simplifyUnOp _ _ (UnOp op (Constant v)) = - constRes =<< doUnOp op v -simplifyUnOp defOf _ (UnOp Not (Var v)) - | Just (BasicOp (UnOp Not v2), v_cs) <- defOf v = - Just (SubExp v2, v_cs) -simplifyUnOp _ _ _ = - Nothing - -simplifyConvOp :: SimpleRule lore -simplifyConvOp _ _ (ConvOp op (Constant v)) = - constRes =<< doConvOp op v -simplifyConvOp _ _ (ConvOp op se) - | (from, to) <- convOpType op, - from == to = - subExpRes se -simplifyConvOp lookupVar _ (ConvOp (SExt t2 t1) (Var v)) - | Just (BasicOp (ConvOp (SExt t3 _) se), v_cs) <- lookupVar v, - t2 >= t3 = - Just (ConvOp (SExt t3 t1) se, v_cs) -simplifyConvOp lookupVar _ (ConvOp (ZExt t2 t1) (Var v)) - | Just (BasicOp (ConvOp (ZExt t3 _) se), v_cs) <- lookupVar v, - t2 >= t3 = - Just (ConvOp (ZExt t3 t1) se, v_cs) -simplifyConvOp lookupVar _ (ConvOp (SIToFP t2 t1) (Var v)) - | Just (BasicOp (ConvOp (SExt t3 _) se), v_cs) <- lookupVar v, - t2 >= t3 = - Just (ConvOp (SIToFP t3 t1) se, v_cs) -simplifyConvOp lookupVar _ (ConvOp (UIToFP t2 t1) (Var v)) - | Just (BasicOp (ConvOp (ZExt t3 _) se), v_cs) <- lookupVar v, - t2 >= t3 = - Just (ConvOp (UIToFP t3 t1) se, v_cs) -simplifyConvOp lookupVar _ (ConvOp (FPConv t2 t1) (Var v)) - | Just (BasicOp (ConvOp (FPConv t3 _) se), v_cs) <- lookupVar v, - t2 >= t3 = - Just (ConvOp (FPConv t3 t1) se, v_cs) -simplifyConvOp _ _ _ = - Nothing - --- If expression is true then just replace assertion. -simplifyAssert :: SimpleRule lore -simplifyAssert _ _ (Assert (Constant (BoolValue True)) _ _) = - constRes Checked -simplifyAssert _ _ _ = - Nothing - -constantFoldPrimFun :: BinderOps lore => TopDownRuleGeneric lore -constantFoldPrimFun _ (Let pat (StmAux cs attrs _) (Apply fname args _ _)) - | Just args' <- mapM (isConst . fst) args, - Just (_, _, fun) <- M.lookup (nameToString fname) primFuns, - Just result <- fun args' = - Simplify $ - certifying cs $ - attributing attrs $ - letBind pat $ BasicOp $ SubExp $ Constant result - where - isConst (Constant v) = Just v - isConst _ = Nothing -constantFoldPrimFun _ _ = Skip - -simplifyIndex :: BinderOps lore => BottomUpRuleBasicOp lore -simplifyIndex (vtable, used) pat@(Pattern [] [pe]) (StmAux cs attrs _) (Index idd inds) - | Just m <- simplifyIndexing vtable seType idd inds consumed = Simplify $ do - res <- m - attributing attrs $ case res of - SubExpResult cs' se -> - certifying (cs <> cs') $ - letBindNames (patternNames pat) $ BasicOp $ SubExp se - IndexResult extra_cs idd' inds' -> - certifying (cs <> extra_cs) $ - letBindNames (patternNames pat) $ BasicOp $ Index idd' inds' - where - consumed = patElemName pe `UT.isConsumed` used - seType (Var v) = ST.lookupType v vtable - seType (Constant v) = Just $ Prim $ primValueType v -simplifyIndex _ _ _ _ = Skip - -data IndexResult - = IndexResult Certificates VName (Slice SubExp) - | SubExpResult Certificates SubExp - -simplifyIndexing :: - MonadBinder m => - ST.SymbolTable (Lore m) -> - TypeLookup -> - VName -> - Slice SubExp -> - Bool -> - Maybe (m IndexResult) -simplifyIndexing vtable seType idd inds consuming = - case defOf idd of - _ - | Just t <- seType (Var idd), - inds == fullSlice t [] -> - Just $ pure $ SubExpResult mempty $ Var idd - | Just inds' <- sliceIndices inds, - Just (ST.Indexed cs e) <- ST.index idd inds' vtable, - worthInlining e, - all (`ST.elem` vtable) (unCertificates cs) -> - Just $ SubExpResult cs <$> toSubExp "index_primexp" e - | Just inds' <- sliceIndices inds, - Just (ST.IndexedArray cs arr inds'') <- ST.index idd inds' vtable, - all (worthInlining . untyped) inds'', - all (`ST.elem` vtable) (unCertificates cs) -> - Just $ - IndexResult cs arr . map DimFix - <$> mapM (toSubExp "index_primexp") inds'' - Nothing -> Nothing - Just (SubExp (Var v), cs) -> Just $ pure $ IndexResult cs v inds - Just (Iota _ x s to_it, cs) - | [DimFix ii] <- inds, - Just (Prim (IntType from_it)) <- seType ii -> - Just $ - let mul = BinOpExp $ Mul to_it OverflowWrap - add = BinOpExp $ Add to_it OverflowWrap - in fmap (SubExpResult cs) $ - toSubExp "index_iota" $ - ( sExt to_it (primExpFromSubExp (IntType from_it) ii) - `mul` primExpFromSubExp (IntType to_it) s - ) - `add` primExpFromSubExp (IntType to_it) x - | [DimSlice i_offset i_n i_stride] <- inds -> - Just $ do - i_offset' <- asIntS to_it i_offset - i_stride' <- asIntS to_it i_stride - let mul = BinOpExp $ Mul to_it OverflowWrap - add = BinOpExp $ Add to_it OverflowWrap - i_offset'' <- - toSubExp "iota_offset" $ - ( primExpFromSubExp (IntType to_it) x - `mul` primExpFromSubExp (IntType to_it) s - ) - `add` primExpFromSubExp (IntType to_it) i_offset' - i_stride'' <- - letSubExp "iota_offset" $ - BasicOp $ BinOp (Mul Int64 OverflowWrap) s i_stride' - fmap (SubExpResult cs) $ - letSubExp "slice_iota" $ - BasicOp $ Iota i_n i_offset'' i_stride'' to_it - - -- A rotate cannot be simplified away if we are slicing a rotated dimension. - Just (Rotate offsets a, cs) - | not $ or $ zipWith rotateAndSlice offsets inds -> Just $ do - dims <- arrayDims <$> lookupType a - let adjustI i o d = do - i_p_o <- letSubExp "i_p_o" $ BasicOp $ BinOp (Add Int64 OverflowWrap) i o - letSubExp "rot_i" (BasicOp $ BinOp (SMod Int64 Unsafe) i_p_o d) - adjust (DimFix i, o, d) = - DimFix <$> adjustI i o d - adjust (DimSlice i n s, o, d) = - DimSlice <$> adjustI i o d <*> pure n <*> pure s - IndexResult cs a <$> mapM adjust (zip3 inds offsets dims) - where - rotateAndSlice r DimSlice {} = not $ isCt0 r - rotateAndSlice _ _ = False - Just (Index aa ais, cs) -> - Just $ - IndexResult cs aa - <$> subExpSlice (sliceSlice (primExpSlice ais) (primExpSlice inds)) - Just (Replicate (Shape [_]) (Var vv), cs) - | [DimFix {}] <- inds, not consuming -> Just $ pure $ SubExpResult cs $ Var vv - | DimFix {} : is' <- inds, not consuming -> Just $ pure $ IndexResult cs vv is' - Just (Replicate (Shape [_]) val@(Constant _), cs) - | [DimFix {}] <- inds, not consuming -> Just $ pure $ SubExpResult cs val - Just (Replicate (Shape ds) v, cs) - | (ds_inds, rest_inds) <- splitAt (length ds) inds, - (ds', ds_inds') <- unzip $ mapMaybe index ds_inds, - ds' /= ds -> - Just $ do - arr <- letExp "smaller_replicate" $ BasicOp $ Replicate (Shape ds') v - return $ IndexResult cs arr $ ds_inds' ++ rest_inds - where - index DimFix {} = Nothing - index (DimSlice _ n s) = Just (n, DimSlice (constant (0 :: Int64)) n s) - Just (Rearrange perm src, cs) - | rearrangeReach perm <= length (takeWhile isIndex inds) -> - let inds' = rearrangeShape (rearrangeInverse perm) inds - in Just $ pure $ IndexResult cs src inds' - where - isIndex DimFix {} = True - isIndex _ = False - Just (Copy src, cs) - | Just dims <- arrayDims <$> seType (Var src), - length inds == length dims, - -- It is generally not safe to simplify a slice of a copy, - -- because the result may be used in an in-place update of the - -- original. - Just _ <- mapM dimFix inds, - not consuming, - ST.available src vtable -> - Just $ pure $ IndexResult cs src inds - Just (Reshape newshape src, cs) - | Just newdims <- shapeCoercion newshape, - Just olddims <- arrayDims <$> seType (Var src), - changed_dims <- zipWith (/=) newdims olddims, - not $ or $ drop (length inds) changed_dims -> - Just $ pure $ IndexResult cs src inds - | Just newdims <- shapeCoercion newshape, - Just olddims <- arrayDims <$> seType (Var src), - length newshape == length inds, - length olddims == length newdims -> - Just $ pure $ IndexResult cs src inds - Just (Reshape [_] v2, cs) - | Just [_] <- arrayDims <$> seType (Var v2) -> - Just $ pure $ IndexResult cs v2 inds - Just (Concat d x xs _, cs) - | -- HACK: simplifying the indexing of an N-array concatenation - -- is going to produce an N-deep if expression, which is bad - -- when N is large. To try to avoid that, we use the - -- heuristic not to simplify as long as any of the operands - -- are themselves Concats. The hops it that this will give - -- simplification some time to cut down the concatenation to - -- something smaller, before we start inlining. - not $ any isConcat $ x : xs, - Just (ibef, DimFix i, iaft) <- focusNth d inds, - Just (Prim res_t) <- - (`setArrayDims` sliceDims inds) - <$> ST.lookupType x vtable -> Just $ do - x_len <- arraySize d <$> lookupType x - xs_lens <- mapM (fmap (arraySize d) . lookupType) xs - - let add n m = do - added <- letSubExp "index_concat_add" $ BasicOp $ BinOp (Add Int64 OverflowWrap) n m - return (added, n) - (_, starts) <- mapAccumLM add x_len xs_lens - let xs_and_starts = reverse $ zip xs starts - - let mkBranch [] = - letSubExp "index_concat" $ BasicOp $ Index x $ ibef ++ DimFix i : iaft - mkBranch ((x', start) : xs_and_starts') = do - cmp <- letSubExp "index_concat_cmp" $ BasicOp $ CmpOp (CmpSle Int64) start i - (thisres, thisbnds) <- collectStms $ do - i' <- letSubExp "index_concat_i" $ BasicOp $ BinOp (Sub Int64 OverflowWrap) i start - letSubExp "index_concat" $ BasicOp $ Index x' $ ibef ++ DimFix i' : iaft - thisbody <- mkBodyM thisbnds [thisres] - (altres, altbnds) <- collectStms $ mkBranch xs_and_starts' - altbody <- mkBodyM altbnds [altres] - letSubExp "index_concat_branch" $ - If cmp thisbody altbody $ - IfDec [primBodyType res_t] IfNormal - SubExpResult cs <$> mkBranch xs_and_starts - Just (ArrayLit ses _, cs) - | DimFix (Constant (IntValue (Int64Value i))) : inds' <- inds, - Just se <- maybeNth i ses -> - case inds' of - [] -> Just $ pure $ SubExpResult cs se - _ | Var v2 <- se -> Just $ pure $ IndexResult cs v2 inds' - _ -> Nothing - -- Indexing single-element arrays. We know the index must be 0. - _ - | Just t <- seType $ Var idd, - isCt1 $ arraySize 0 t, - DimFix i : inds' <- inds, - not $ isCt0 i -> - Just $ - pure $ - IndexResult mempty idd $ - DimFix (constant (0 :: Int64)) : inds' - _ -> Nothing - where - defOf v = do - (BasicOp op, def_cs) <- ST.lookupExp v vtable - return (op, def_cs) - worthInlining e - | primExpSizeAtLeast 20 e = False -- totally ad-hoc. - | otherwise = worthInlining' e - worthInlining' (BinOpExp Pow {} _ _) = False - worthInlining' (BinOpExp FPow {} _ _) = False - worthInlining' (BinOpExp _ x y) = worthInlining' x && worthInlining' y - worthInlining' (CmpOpExp _ x y) = worthInlining' x && worthInlining' y - worthInlining' (ConvOpExp _ x) = worthInlining' x - worthInlining' (UnOpExp _ x) = worthInlining' x - worthInlining' FunExp {} = False - worthInlining' _ = True - - isConcat v - | Just (Concat {}, _) <- defOf v = - True - | otherwise = - False - -data ConcatArg - = ArgArrayLit [SubExp] - | ArgReplicate [SubExp] SubExp - | ArgVar VName - -toConcatArg :: ST.SymbolTable lore -> VName -> (ConcatArg, Certificates) -toConcatArg vtable v = - case ST.lookupBasicOp v vtable of - Just (ArrayLit ses _, cs) -> - (ArgArrayLit ses, cs) - Just (Replicate shape se, cs) -> - (ArgReplicate [shapeSize 0 shape] se, cs) - _ -> - (ArgVar v, mempty) - -fromConcatArg :: - MonadBinder m => - Type -> - (ConcatArg, Certificates) -> - m VName -fromConcatArg t (ArgArrayLit ses, cs) = - certifying cs $ letExp "concat_lit" $ BasicOp $ ArrayLit ses $ rowType t -fromConcatArg elem_type (ArgReplicate ws se, cs) = do - let elem_shape = arrayShape elem_type - certifying cs $ do - w <- letSubExp "concat_rep_w" =<< toExp (sum $ map pe64 ws) - letExp "concat_rep" $ BasicOp $ Replicate (setDim 0 elem_shape w) se -fromConcatArg _ (ArgVar v, _) = - pure v - -fuseConcatArg :: - [(ConcatArg, Certificates)] -> - (ConcatArg, Certificates) -> - [(ConcatArg, Certificates)] -fuseConcatArg xs (ArgArrayLit [], _) = - xs -fuseConcatArg xs (ArgReplicate [w] se, cs) - | isCt0 w = - xs - | isCt1 w = - fuseConcatArg xs (ArgArrayLit [se], cs) -fuseConcatArg ((ArgArrayLit x_ses, x_cs) : xs) (ArgArrayLit y_ses, y_cs) = - (ArgArrayLit (x_ses ++ y_ses), x_cs <> y_cs) : xs -fuseConcatArg ((ArgReplicate x_ws x_se, x_cs) : xs) (ArgReplicate y_ws y_se, y_cs) - | x_se == y_se = - (ArgReplicate (x_ws ++ y_ws) x_se, x_cs <> y_cs) : xs -fuseConcatArg xs y = - y : xs - -simplifyConcat :: BinderOps lore => BottomUpRuleBasicOp lore --- concat@1(transpose(x),transpose(y)) == transpose(concat@0(x,y)) -simplifyConcat (vtable, _) pat _ (Concat i x xs new_d) - | Just r <- arrayRank <$> ST.lookupType x vtable, - let perm = [i] ++ [0 .. i -1] ++ [i + 1 .. r -1], - Just (x', x_cs) <- transposedBy perm x, - Just (xs', xs_cs) <- unzip <$> mapM (transposedBy perm) xs = Simplify $ do - concat_rearrange <- - certifying (x_cs <> mconcat xs_cs) $ - letExp "concat_rearrange" $ BasicOp $ Concat 0 x' xs' new_d - letBind pat $ BasicOp $ Rearrange perm concat_rearrange - where - transposedBy perm1 v = - case ST.lookupExp v vtable of - Just (BasicOp (Rearrange perm2 v'), vcs) - | perm1 == perm2 -> Just (v', vcs) - _ -> Nothing - --- Removing a concatenation that involves only a single array. This --- may be produced as a result of other simplification rules. -simplifyConcat _ pat aux (Concat _ x [] _) = - Simplify $ - -- Still need a copy because Concat produces a fresh array. - auxing aux $ letBind pat $ BasicOp $ Copy x --- concat xs (concat ys zs) == concat xs ys zs -simplifyConcat (vtable, _) pat (StmAux cs attrs _) (Concat i x xs new_d) - | x' /= x || concat xs' /= xs = - Simplify $ - certifying (cs <> x_cs <> mconcat xs_cs) $ - attributing attrs $ - letBind pat $ - BasicOp $ Concat i x' (zs ++ concat xs') new_d - where - (x' : zs, x_cs) = isConcat x - (xs', xs_cs) = unzip $ map isConcat xs - isConcat v = case ST.lookupBasicOp v vtable of - Just (Concat j y ys _, v_cs) | j == i -> (y : ys, v_cs) - _ -> ([v], mempty) - --- Fusing arguments to the concat when possible. Only done when --- concatenating along the outer dimension for now. -simplifyConcat (vtable, _) pat aux (Concat 0 x xs outer_w) - | -- We produce the to-be-concatenated arrays in reverse order, so - -- reverse them back. - y : ys <- - reverse $ - foldl' fuseConcatArg mempty $ - map (toConcatArg vtable) $ x : xs, - length xs /= length ys = - Simplify $ do - elem_type <- lookupType x - y' <- fromConcatArg elem_type y - ys' <- mapM (fromConcatArg elem_type) ys - auxing aux $ letBind pat $ BasicOp $ Concat 0 y' ys' outer_w -simplifyConcat _ _ _ _ = Skip - -ruleIf :: BinderOps lore => TopDownRuleIf lore -ruleIf _ pat _ (e1, tb, fb, IfDec _ ifsort) - | Just branch <- checkBranch, - ifsort /= IfFallback || isCt1 e1 = Simplify $ do - let ses = bodyResult branch - addStms $ bodyStms branch - sequence_ - [ letBindNames [patElemName p] $ BasicOp $ SubExp se - | (p, se) <- zip (patternElements pat) ses - ] - where - checkBranch - | isCt1 e1 = Just tb - | isCt0 e1 = Just fb - | otherwise = Nothing - --- IMPROVE: the following two rules can be generalised to work in more --- cases, especially when the branches have bindings, or return more --- than one value. --- --- if c then True else v == c || v -ruleIf - _ - pat - _ - ( cond, - Body _ tstms [Constant (BoolValue True)], - Body _ fstms [se], - IfDec ts _ - ) - | null tstms, - null fstms, - [Prim Bool] <- map extTypeOf ts = - Simplify $ letBind pat $ BasicOp $ BinOp LogOr cond se --- When type(x)==bool, if c then x else y == (c && x) || (!c && y) -ruleIf _ pat _ (cond, tb, fb, IfDec ts _) - | Body _ tstms [tres] <- tb, - Body _ fstms [fres] <- fb, - all (safeExp . stmExp) $ tstms <> fstms, - all ((== Prim Bool) . extTypeOf) ts = Simplify $ do - addStms tstms - addStms fstms - e <- - eBinOp - LogOr - (pure $ BasicOp $ BinOp LogAnd cond tres) - ( eBinOp - LogAnd - (pure $ BasicOp $ UnOp Not cond) - (pure $ BasicOp $ SubExp fres) - ) - letBind pat e -ruleIf _ pat _ (_, tbranch, _, IfDec _ IfFallback) - | null $ patternContextNames pat, - all (safeExp . stmExp) $ bodyStms tbranch = Simplify $ do - let ses = bodyResult tbranch - addStms $ bodyStms tbranch - sequence_ - [ letBindNames [patElemName p] $ BasicOp $ SubExp se - | (p, se) <- zip (patternElements pat) ses - ] -ruleIf _ pat _ (cond, tb, fb, _) - | Body _ _ [Constant (IntValue t)] <- tb, - Body _ _ [Constant (IntValue f)] <- fb = - if oneIshInt t && zeroIshInt f - then - Simplify $ - letBind pat $ BasicOp $ ConvOp (BToI (intValueType t)) cond - else - if zeroIshInt t && oneIshInt f - then Simplify $ do - cond_neg <- letSubExp "cond_neg" $ BasicOp $ UnOp Not cond - letBind pat $ BasicOp $ ConvOp (BToI (intValueType t)) cond_neg - else Skip -ruleIf _ _ _ _ = Skip - --- | Move out results of a conditional expression whose computation is --- either invariant to the branches (only done for results in the --- context), or the same in both branches. -hoistBranchInvariant :: BinderOps lore => TopDownRuleIf lore -hoistBranchInvariant _ pat _ (cond, tb, fb, IfDec ret ifsort) = Simplify $ do - let tses = bodyResult tb - fses = bodyResult fb - (hoistings, (pes, ts, res)) <- - fmap (fmap unzip3 . partitionEithers) $ - mapM branchInvariant $ - zip3 - (patternElements pat) - (map Left [0 .. num_ctx -1] ++ map Right ret) - (zip tses fses) - let ctx_fixes = catMaybes hoistings - (tses', fses') = unzip res - tb' = tb {bodyResult = tses'} - fb' = fb {bodyResult = fses'} - ret' = foldr (uncurry fixExt) (rights ts) ctx_fixes - (ctx_pes, val_pes) = splitFromEnd (length ret') pes - if not $ null hoistings -- Was something hoisted? - then do - -- We may have to add some reshapes if we made the type - -- less existential. - tb'' <- reshapeBodyResults tb' $ map extTypeOf ret' - fb'' <- reshapeBodyResults fb' $ map extTypeOf ret' - letBind (Pattern ctx_pes val_pes) $ - If cond tb'' fb'' (IfDec ret' ifsort) - else cannotSimplify - where - num_ctx = length $ patternContextElements pat - bound_in_branches = - namesFromList $ - concatMap (patternNames . stmPattern) $ - bodyStms tb <> bodyStms fb - mem_sizes = freeIn $ filter (isMem . patElemType) $ patternElements pat - invariant Constant {} = True - invariant (Var v) = not $ v `nameIn` bound_in_branches - - isMem Mem {} = True - isMem _ = False - sizeOfMem v = v `nameIn` mem_sizes - - branchInvariant (pe, t, (tse, fse)) - -- Do both branches return the same value? - | tse == fse = do - letBindNames [patElemName pe] $ BasicOp $ SubExp tse - hoisted pe t - - -- Do both branches return values that are free in the - -- branch, and are we not the only pattern element? The - -- latter is to avoid infinite application of this rule. - | invariant tse, - invariant fse, - patternSize pat > 1, - Prim _ <- patElemType pe, - not $ sizeOfMem $ patElemName pe = do - bt <- expTypesFromPattern $ Pattern [] [pe] - letBindNames [patElemName pe] - =<< ( If cond <$> resultBodyM [tse] - <*> resultBodyM [fse] - <*> pure (IfDec bt ifsort) - ) - hoisted pe t - | otherwise = - return $ Right (pe, t, (tse, fse)) - - hoisted pe (Left i) = return $ Left $ Just (i, Var $ patElemName pe) - hoisted _ Right {} = return $ Left Nothing - - reshapeBodyResults body rets = insertStmsM $ do - ses <- bodyBind body - let (ctx_ses, val_ses) = splitFromEnd (length rets) ses - resultBodyM . (ctx_ses ++) =<< zipWithM reshapeResult val_ses rets - reshapeResult (Var v) t@Array {} = do - v_t <- lookupType v - let newshape = arrayDims $ removeExistentials t v_t - if newshape /= arrayDims v_t - then letSubExp "branch_ctx_reshaped" $ shapeCoerce newshape v - else return $ Var v - reshapeResult se _ = - return se - -simplifyIdentityReshape :: SimpleRule lore -simplifyIdentityReshape _ seType (Reshape newshape v) - | Just t <- seType $ Var v, - newDims newshape == arrayDims t -- No-op reshape. - = - subExpRes $ Var v -simplifyIdentityReshape _ _ _ = Nothing - -simplifyReshapeReshape :: SimpleRule lore -simplifyReshapeReshape defOf _ (Reshape newshape v) - | Just (BasicOp (Reshape oldshape v2), v_cs) <- defOf v = - Just (Reshape (fuseReshape oldshape newshape) v2, v_cs) -simplifyReshapeReshape _ _ _ = Nothing - -simplifyReshapeScratch :: SimpleRule lore -simplifyReshapeScratch defOf _ (Reshape newshape v) - | Just (BasicOp (Scratch bt _), v_cs) <- defOf v = - Just (Scratch bt $ newDims newshape, v_cs) -simplifyReshapeScratch _ _ _ = Nothing - -simplifyReshapeReplicate :: SimpleRule lore -simplifyReshapeReplicate defOf seType (Reshape newshape v) - | Just (BasicOp (Replicate _ se), v_cs) <- defOf v, - Just oldshape <- arrayShape <$> seType se, - shapeDims oldshape `isSuffixOf` newDims newshape = - let new = - take (length newshape - shapeRank oldshape) $ - newDims newshape - in Just (Replicate (Shape new) se, v_cs) -simplifyReshapeReplicate _ _ _ = Nothing - -simplifyReshapeIota :: SimpleRule lore -simplifyReshapeIota defOf _ (Reshape newshape v) - | Just (BasicOp (Iota _ offset stride it), v_cs) <- defOf v, - [n] <- newDims newshape = - Just (Iota n offset stride it, v_cs) -simplifyReshapeIota _ _ _ = Nothing - -improveReshape :: SimpleRule lore -improveReshape _ seType (Reshape newshape v) - | Just t <- seType $ Var v, - newshape' <- informReshape (arrayDims t) newshape, - newshape' /= newshape = - Just (Reshape newshape' v, mempty) -improveReshape _ _ _ = Nothing - --- | If we are copying a scratch array (possibly indirectly), just turn it into a scratch by --- itself. -copyScratchToScratch :: SimpleRule lore -copyScratchToScratch defOf seType (Copy src) = do - t <- seType $ Var src - if isActuallyScratch src - then Just (Scratch (elemType t) (arrayDims t), mempty) - else Nothing - where - isActuallyScratch v = - case asBasicOp . fst =<< defOf v of - Just Scratch {} -> True - Just (Rearrange _ v') -> isActuallyScratch v' - Just (Reshape _ v') -> isActuallyScratch v' - _ -> False -copyScratchToScratch _ _ _ = - Nothing - -ruleBasicOp :: BinderOps lore => TopDownRuleBasicOp lore --- Check all the simpleRules. -ruleBasicOp vtable pat aux op - | Just (op', cs) <- msum [rule defOf seType op | rule <- simpleRules] = - Simplify $ certifying (cs <> stmAuxCerts aux) $ letBind pat $ BasicOp op' - where - defOf = (`ST.lookupExp` vtable) - seType (Var v) = ST.lookupType v vtable - seType (Constant v) = Just $ Prim $ primValueType v -ruleBasicOp vtable pat _ (Update src _ (Var v)) - | Just (BasicOp Scratch {}, _) <- ST.lookupExp v vtable = - Simplify $ letBind pat $ BasicOp $ SubExp $ Var src --- If we are writing a single-element slice from some array, and the --- element of that array can be computed as a PrimExp based on the --- index, let's just write that instead. -ruleBasicOp vtable pat aux (Update src [DimSlice i n s] (Var v)) - | isCt1 n, - isCt1 s, - Just (ST.Indexed cs e) <- ST.index v [intConst Int64 0] vtable = - Simplify $ do - e' <- toSubExp "update_elem" e - auxing aux $ - certifying cs $ - letBind pat $ BasicOp $ Update src [DimFix i] e' -ruleBasicOp vtable pat _ (Update dest destis (Var v)) - | Just (e, _) <- ST.lookupExp v vtable, - arrayFrom e = - Simplify $ letBind pat $ BasicOp $ SubExp $ Var dest - where - arrayFrom (BasicOp (Copy copy_v)) - | Just (e', _) <- ST.lookupExp copy_v vtable = - arrayFrom e' - arrayFrom (BasicOp (Index src srcis)) = - src == dest && destis == srcis - arrayFrom (BasicOp (Replicate v_shape v_se)) - | Just (Replicate dest_shape dest_se, _) <- ST.lookupBasicOp dest vtable, - v_se == dest_se, - shapeDims v_shape `isSuffixOf` shapeDims dest_shape = - True - arrayFrom _ = - False -ruleBasicOp vtable pat _ (Update dest is se) - | Just dest_t <- ST.lookupType dest vtable, - isFullSlice (arrayShape dest_t) is = Simplify $ - case se of - Var v | not $ null $ sliceDims is -> do - v_reshaped <- - letExp (baseString v ++ "_reshaped") $ - BasicOp $ Reshape (map DimNew $ arrayDims dest_t) v - letBind pat $ BasicOp $ Copy v_reshaped - _ -> letBind pat $ BasicOp $ ArrayLit [se] $ rowType dest_t -ruleBasicOp vtable pat (StmAux cs1 attrs _) (Update dest1 is1 (Var v1)) - | Just (Update dest2 is2 se2, cs2) <- ST.lookupBasicOp v1 vtable, - Just (Copy v3, cs3) <- ST.lookupBasicOp dest2 vtable, - Just (Index v4 is4, cs4) <- ST.lookupBasicOp v3 vtable, - is4 == is1, - v4 == dest1 = - Simplify $ - certifying (cs1 <> cs2 <> cs3 <> cs4) $ do - is5 <- subExpSlice $ sliceSlice (primExpSlice is1) (primExpSlice is2) - attributing attrs $ letBind pat $ BasicOp $ Update dest1 is5 se2 -ruleBasicOp vtable pat _ (CmpOp (CmpEq t) se1 se2) - | Just m <- simplifyWith se1 se2 = Simplify m - | Just m <- simplifyWith se2 se1 = Simplify m - where - simplifyWith (Var v) x - | Just bnd <- ST.lookupStm v vtable, - If p tbranch fbranch _ <- stmExp bnd, - Just (y, z) <- - returns v (stmPattern bnd) tbranch fbranch, - not $ boundInBody tbranch `namesIntersect` freeIn y, - not $ boundInBody fbranch `namesIntersect` freeIn z = Just $ do - eq_x_y <- - letSubExp "eq_x_y" $ BasicOp $ CmpOp (CmpEq t) x y - eq_x_z <- - letSubExp "eq_x_z" $ BasicOp $ CmpOp (CmpEq t) x z - p_and_eq_x_y <- - letSubExp "p_and_eq_x_y" $ BasicOp $ BinOp LogAnd p eq_x_y - not_p <- - letSubExp "not_p" $ BasicOp $ UnOp Not p - not_p_and_eq_x_z <- - letSubExp "p_and_eq_x_y" $ BasicOp $ BinOp LogAnd not_p eq_x_z - letBind pat $ - BasicOp $ BinOp LogOr p_and_eq_x_y not_p_and_eq_x_z - simplifyWith _ _ = - Nothing - - returns v ifpat tbranch fbranch = - fmap snd $ - find ((== v) . patElemName . fst) $ - zip (patternValueElements ifpat) $ - zip (bodyResult tbranch) (bodyResult fbranch) -ruleBasicOp _ pat _ (Replicate (Shape []) se@Constant {}) = - Simplify $ letBind pat $ BasicOp $ SubExp se -ruleBasicOp _ pat _ (Replicate (Shape []) (Var v)) = Simplify $ do - v_t <- lookupType v - letBind pat $ - BasicOp $ - if primType v_t - then SubExp $ Var v - else Copy v -ruleBasicOp vtable pat _ (Replicate shape (Var v)) - | Just (BasicOp (Replicate shape2 se), cs) <- ST.lookupExp v vtable = - Simplify $ certifying cs $ letBind pat $ BasicOp $ Replicate (shape <> shape2) se -ruleBasicOp _ pat _ (ArrayLit (se : ses) _) - | all (== se) ses = - Simplify $ - let n = constant (fromIntegral (length ses) + 1 :: Int64) - in letBind pat $ BasicOp $ Replicate (Shape [n]) se -ruleBasicOp vtable pat aux (Index idd slice) - | Just inds <- sliceIndices slice, - Just (BasicOp (Reshape newshape idd2), idd_cs) <- ST.lookupExp idd vtable, - length newshape == length inds = - Simplify $ - case shapeCoercion newshape of - Just _ -> - certifying idd_cs $ - auxing aux $ - letBind pat $ BasicOp $ Index idd2 slice - Nothing -> do - -- Linearise indices and map to old index space. - oldshape <- arrayDims <$> lookupType idd2 - let new_inds = - reshapeIndex - (map pe64 oldshape) - (map pe64 $ newDims newshape) - (map pe64 inds) - new_inds' <- - mapM (toSubExp "new_index") new_inds - certifying idd_cs $ - auxing aux $ - letBind pat $ BasicOp $ Index idd2 $ map DimFix new_inds' -ruleBasicOp _ pat _ (BinOp (Pow t) e1 e2) - | e1 == intConst t 2 = - Simplify $ letBind pat $ BasicOp $ BinOp (Shl t) (intConst t 1) e2 --- Handle identity permutation. -ruleBasicOp _ pat _ (Rearrange perm v) - | sort perm == perm = - Simplify $ letBind pat $ BasicOp $ SubExp $ Var v -ruleBasicOp vtable pat aux (Rearrange perm v) - | Just (BasicOp (Rearrange perm2 e), v_cs) <- ST.lookupExp v vtable = - -- Rearranging a rearranging: compose the permutations. - Simplify $ - certifying v_cs $ - auxing aux $ - letBind pat $ BasicOp $ Rearrange (perm `rearrangeCompose` perm2) e -ruleBasicOp vtable pat aux (Rearrange perm v) - | Just (BasicOp (Rotate offsets v2), v_cs) <- ST.lookupExp v vtable, - Just (BasicOp (Rearrange perm3 v3), v2_cs) <- ST.lookupExp v2 vtable = Simplify $ do - let offsets' = rearrangeShape (rearrangeInverse perm3) offsets - rearrange_rotate <- letExp "rearrange_rotate" $ BasicOp $ Rotate offsets' v3 - certifying (v_cs <> v2_cs) $ - auxing aux $ - letBind pat $ BasicOp $ Rearrange (perm `rearrangeCompose` perm3) rearrange_rotate - --- Rearranging a replicate where the outer dimension is left untouched. -ruleBasicOp vtable pat aux (Rearrange perm v1) - | Just (BasicOp (Replicate dims (Var v2)), v1_cs) <- ST.lookupExp v1 vtable, - num_dims <- shapeRank dims, - (rep_perm, rest_perm) <- splitAt num_dims perm, - not $ null rest_perm, - rep_perm == [0 .. length rep_perm -1] = - Simplify $ - certifying v1_cs $ - auxing aux $ do - v <- - letSubExp "rearrange_replicate" $ - BasicOp $ Rearrange (map (subtract num_dims) rest_perm) v2 - letBind pat $ BasicOp $ Replicate dims v - --- A zero-rotation is identity. -ruleBasicOp _ pat _ (Rotate offsets v) - | all isCt0 offsets = Simplify $ letBind pat $ BasicOp $ SubExp $ Var v -ruleBasicOp vtable pat aux (Rotate offsets v) - | Just (BasicOp (Rearrange perm v2), v_cs) <- ST.lookupExp v vtable, - Just (BasicOp (Rotate offsets2 v3), v2_cs) <- ST.lookupExp v2 vtable = Simplify $ do - let offsets2' = rearrangeShape (rearrangeInverse perm) offsets2 - addOffsets x y = letSubExp "summed_offset" $ BasicOp $ BinOp (Add Int64 OverflowWrap) x y - offsets' <- zipWithM addOffsets offsets offsets2' - rotate_rearrange <- - auxing aux $ letExp "rotate_rearrange" $ BasicOp $ Rearrange perm v3 - certifying (v_cs <> v2_cs) $ - letBind pat $ BasicOp $ Rotate offsets' rotate_rearrange - --- Combining Rotates. -ruleBasicOp vtable pat aux (Rotate offsets1 v) - | Just (BasicOp (Rotate offsets2 v2), v_cs) <- ST.lookupExp v vtable = Simplify $ do - offsets <- zipWithM add offsets1 offsets2 - certifying v_cs $ - auxing aux $ - letBind pat $ BasicOp $ Rotate offsets v2 - where - add x y = letSubExp "offset" $ BasicOp $ BinOp (Add Int64 OverflowWrap) x y - --- If we see an Update with a scalar where the value to be written is --- the result of indexing some other array, then we convert it into an --- Update with a slice of that array. This matters when the arrays --- are far away (on the GPU, say), because it avoids a copy of the --- scalar to and from the host. -ruleBasicOp vtable pat aux (Update arr_x slice_x (Var v)) - | Just _ <- sliceIndices slice_x, - Just (Index arr_y slice_y, cs_y) <- ST.lookupBasicOp v vtable, - ST.available arr_y vtable, - -- XXX: we should check for proper aliasing here instead. - arr_y /= arr_x, - Just (slice_x_bef, DimFix i, []) <- focusNth (length slice_x - 1) slice_x, - Just (slice_y_bef, DimFix j, []) <- focusNth (length slice_y - 1) slice_y = Simplify $ do - let slice_x' = slice_x_bef ++ [DimSlice i (intConst Int64 1) (intConst Int64 1)] - slice_y' = slice_y_bef ++ [DimSlice j (intConst Int64 1) (intConst Int64 1)] - v' <- letExp (baseString v ++ "_slice") $ BasicOp $ Index arr_y slice_y' - certifying cs_y $ - auxing aux $ - letBind pat $ BasicOp $ Update arr_x slice_x' $ Var v' - --- Simplify away 0<=i when 'i' is from a loop of form 'for i < n'. -ruleBasicOp vtable pat aux (CmpOp CmpSle {} x y) - | Constant (IntValue (Int64Value 0)) <- x, - Var v <- y, - Just _ <- ST.lookupLoopVar v vtable = - Simplify $ auxing aux $ letBind pat $ BasicOp $ SubExp $ constant True --- Simplify away i [TopDownRule lore] +topDownRules = + [ RuleGeneric constantFoldPrimFun, + RuleIf ruleIf, + RuleIf hoistBranchInvariant + ] + +bottomUpRules :: BinderOps lore => [BottomUpRule lore] +bottomUpRules = + [ RuleIf removeDeadBranchResult, + RuleBasicOp simplifyIndex + ] + +-- | A set of standard simplification rules. These assume pure +-- functional semantics, and so probably should not be applied after +-- memory block merging. +standardRules :: (BinderOps lore, Aliased lore) => RuleBook lore +standardRules = ruleBook topDownRules bottomUpRules <> loopRules <> basicOpRules + +-- | Turn @copy(x)@ into @x@ iff @x@ is not used after this copy +-- statement and it can be consumed. +-- +-- This simplistic rule is only valid before we introduce memory. +removeUnnecessaryCopy :: BinderOps lore => BottomUpRuleBasicOp lore +removeUnnecessaryCopy (vtable, used) (Pattern [] [d]) _ (Copy v) + | not (v `UT.isConsumed` used), + (not (v `UT.used` used) && consumable) || not (patElemName d `UT.isConsumed` used) = + Simplify $ letBindNames [patElemName d] $ BasicOp $ SubExp $ Var v + where + -- We need to make sure we can even consume the original. + -- This is currently a hacky check, much too conservative, + -- because we don't have the information conveniently + -- available. + consumable = case M.lookup v $ ST.toScope vtable of + Just (FParamName info) -> unique $ declTypeOf info + _ -> False +removeUnnecessaryCopy _ _ _ _ = Skip + +constantFoldPrimFun :: BinderOps lore => TopDownRuleGeneric lore +constantFoldPrimFun _ (Let pat (StmAux cs attrs _) (Apply fname args _ _)) + | Just args' <- mapM (isConst . fst) args, + Just (_, _, fun) <- M.lookup (nameToString fname) primFuns, + Just result <- fun args' = + Simplify $ + certifying cs $ + attributing attrs $ + letBind pat $ BasicOp $ SubExp $ Constant result + where + isConst (Constant v) = Just v + isConst _ = Nothing +constantFoldPrimFun _ _ = Skip + +simplifyIndex :: BinderOps lore => BottomUpRuleBasicOp lore +simplifyIndex (vtable, used) pat@(Pattern [] [pe]) (StmAux cs attrs _) (Index idd inds) + | Just m <- simplifyIndexing vtable seType idd inds consumed = Simplify $ do + res <- m + attributing attrs $ case res of + SubExpResult cs' se -> + certifying (cs <> cs') $ + letBindNames (patternNames pat) $ BasicOp $ SubExp se + IndexResult extra_cs idd' inds' -> + certifying (cs <> extra_cs) $ + letBindNames (patternNames pat) $ BasicOp $ Index idd' inds' + where + consumed = patElemName pe `UT.isConsumed` used + seType (Var v) = ST.lookupType v vtable + seType (Constant v) = Just $ Prim $ primValueType v +simplifyIndex _ _ _ _ = Skip + +ruleIf :: BinderOps lore => TopDownRuleIf lore +ruleIf _ pat _ (e1, tb, fb, IfDec _ ifsort) + | Just branch <- checkBranch, + ifsort /= IfFallback || isCt1 e1 = Simplify $ do + let ses = bodyResult branch + addStms $ bodyStms branch + sequence_ + [ letBindNames [patElemName p] $ BasicOp $ SubExp se + | (p, se) <- zip (patternElements pat) ses + ] + where + checkBranch + | isCt1 e1 = Just tb + | isCt0 e1 = Just fb + | otherwise = Nothing + +-- IMPROVE: the following two rules can be generalised to work in more +-- cases, especially when the branches have bindings, or return more +-- than one value. +-- +-- if c then True else v == c || v +ruleIf + _ + pat + _ + ( cond, + Body _ tstms [Constant (BoolValue True)], + Body _ fstms [se], + IfDec ts _ + ) + | null tstms, + null fstms, + [Prim Bool] <- map extTypeOf ts = + Simplify $ letBind pat $ BasicOp $ BinOp LogOr cond se +-- When type(x)==bool, if c then x else y == (c && x) || (!c && y) +ruleIf _ pat _ (cond, tb, fb, IfDec ts _) + | Body _ tstms [tres] <- tb, + Body _ fstms [fres] <- fb, + all (safeExp . stmExp) $ tstms <> fstms, + all ((== Prim Bool) . extTypeOf) ts = Simplify $ do + addStms tstms + addStms fstms + e <- + eBinOp + LogOr + (pure $ BasicOp $ BinOp LogAnd cond tres) + ( eBinOp + LogAnd + (pure $ BasicOp $ UnOp Not cond) + (pure $ BasicOp $ SubExp fres) + ) + letBind pat e +ruleIf _ pat _ (_, tbranch, _, IfDec _ IfFallback) + | null $ patternContextNames pat, + all (safeExp . stmExp) $ bodyStms tbranch = Simplify $ do + let ses = bodyResult tbranch + addStms $ bodyStms tbranch + sequence_ + [ letBindNames [patElemName p] $ BasicOp $ SubExp se + | (p, se) <- zip (patternElements pat) ses + ] +ruleIf _ pat _ (cond, tb, fb, _) + | Body _ _ [Constant (IntValue t)] <- tb, + Body _ _ [Constant (IntValue f)] <- fb = + if oneIshInt t && zeroIshInt f + then + Simplify $ + letBind pat $ BasicOp $ ConvOp (BToI (intValueType t)) cond + else + if zeroIshInt t && oneIshInt f + then Simplify $ do + cond_neg <- letSubExp "cond_neg" $ BasicOp $ UnOp Not cond + letBind pat $ BasicOp $ ConvOp (BToI (intValueType t)) cond_neg + else Skip +ruleIf _ _ _ _ = Skip + +-- | Move out results of a conditional expression whose computation is +-- either invariant to the branches (only done for results in the +-- context), or the same in both branches. +hoistBranchInvariant :: BinderOps lore => TopDownRuleIf lore +hoistBranchInvariant _ pat _ (cond, tb, fb, IfDec ret ifsort) = Simplify $ do + let tses = bodyResult tb + fses = bodyResult fb + (hoistings, (pes, ts, res)) <- + fmap (fmap unzip3 . partitionEithers) $ + mapM branchInvariant $ + zip3 + (patternElements pat) + (map Left [0 .. num_ctx -1] ++ map Right ret) + (zip tses fses) + let ctx_fixes = catMaybes hoistings + (tses', fses') = unzip res + tb' = tb {bodyResult = tses'} + fb' = fb {bodyResult = fses'} + ret' = foldr (uncurry fixExt) (rights ts) ctx_fixes + (ctx_pes, val_pes) = splitFromEnd (length ret') pes + if not $ null hoistings -- Was something hoisted? + then do + -- We may have to add some reshapes if we made the type + -- less existential. + tb'' <- reshapeBodyResults tb' $ map extTypeOf ret' + fb'' <- reshapeBodyResults fb' $ map extTypeOf ret' + letBind (Pattern ctx_pes val_pes) $ + If cond tb'' fb'' (IfDec ret' ifsort) + else cannotSimplify + where + num_ctx = length $ patternContextElements pat + bound_in_branches = + namesFromList $ + concatMap (patternNames . stmPattern) $ + bodyStms tb <> bodyStms fb + mem_sizes = freeIn $ filter (isMem . patElemType) $ patternElements pat + invariant Constant {} = True + invariant (Var v) = not $ v `nameIn` bound_in_branches + + isMem Mem {} = True + isMem _ = False + sizeOfMem v = v `nameIn` mem_sizes + + branchInvariant (pe, t, (tse, fse)) + -- Do both branches return the same value? + | tse == fse = do + letBindNames [patElemName pe] $ BasicOp $ SubExp tse + hoisted pe t + + -- Do both branches return values that are free in the + -- branch, and are we not the only pattern element? The + -- latter is to avoid infinite application of this rule. + | invariant tse, + invariant fse, + patternSize pat > 1, + Prim _ <- patElemType pe, + not $ sizeOfMem $ patElemName pe = do + bt <- expTypesFromPattern $ Pattern [] [pe] + letBindNames [patElemName pe] + =<< ( If cond <$> resultBodyM [tse] + <*> resultBodyM [fse] + <*> pure (IfDec bt ifsort) + ) + hoisted pe t + | otherwise = + return $ Right (pe, t, (tse, fse)) + + hoisted pe (Left i) = return $ Left $ Just (i, Var $ patElemName pe) + hoisted _ Right {} = return $ Left Nothing + + reshapeBodyResults body rets = insertStmsM $ do + ses <- bodyBind body + let (ctx_ses, val_ses) = splitFromEnd (length rets) ses + resultBodyM . (ctx_ses ++) =<< zipWithM reshapeResult val_ses rets + reshapeResult (Var v) t@Array {} = do + v_t <- lookupType v + let newshape = arrayDims $ removeExistentials t v_t + if newshape /= arrayDims v_t + then letSubExp "branch_ctx_reshaped" $ shapeCoerce newshape v + else return $ Var v + reshapeResult se _ = + return se -- | Remove the return values of a branch, that are not actually used -- after a branch. Standard dead code removal can remove the branch diff --git a/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs b/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/Simplify/Rules/BasicOp.hs @@ -0,0 +1,384 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_GHC -Wno-overlapping-patterns -Wno-incomplete-patterns -Wno-incomplete-uni-patterns -Wno-incomplete-record-updates #-} + +-- | Some simplification rules for 'BasicOp'. +module Futhark.Optimise.Simplify.Rules.BasicOp + ( basicOpRules, + ) +where + +import Control.Monad +import Data.List (find, foldl', isSuffixOf, sort) +import Futhark.Analysis.PrimExp.Convert +import qualified Futhark.Analysis.SymbolTable as ST +import Futhark.Construct +import Futhark.IR +import Futhark.IR.Prop.Aliases +import Futhark.Optimise.Simplify.Rule +import Futhark.Optimise.Simplify.Rules.Loop +import Futhark.Optimise.Simplify.Rules.Simple +import Futhark.Util + +isCt1 :: SubExp -> Bool +isCt1 (Constant v) = oneIsh v +isCt1 _ = False + +isCt0 :: SubExp -> Bool +isCt0 (Constant v) = zeroIsh v +isCt0 _ = False + +data ConcatArg + = ArgArrayLit [SubExp] + | ArgReplicate [SubExp] SubExp + | ArgVar VName + +toConcatArg :: ST.SymbolTable lore -> VName -> (ConcatArg, Certificates) +toConcatArg vtable v = + case ST.lookupBasicOp v vtable of + Just (ArrayLit ses _, cs) -> + (ArgArrayLit ses, cs) + Just (Replicate shape se, cs) -> + (ArgReplicate [shapeSize 0 shape] se, cs) + _ -> + (ArgVar v, mempty) + +fromConcatArg :: + MonadBinder m => + Type -> + (ConcatArg, Certificates) -> + m VName +fromConcatArg t (ArgArrayLit ses, cs) = + certifying cs $ letExp "concat_lit" $ BasicOp $ ArrayLit ses $ rowType t +fromConcatArg elem_type (ArgReplicate ws se, cs) = do + let elem_shape = arrayShape elem_type + certifying cs $ do + w <- letSubExp "concat_rep_w" =<< toExp (sum $ map pe64 ws) + letExp "concat_rep" $ BasicOp $ Replicate (setDim 0 elem_shape w) se +fromConcatArg _ (ArgVar v, _) = + pure v + +fuseConcatArg :: + [(ConcatArg, Certificates)] -> + (ConcatArg, Certificates) -> + [(ConcatArg, Certificates)] +fuseConcatArg xs (ArgArrayLit [], _) = + xs +fuseConcatArg xs (ArgReplicate [w] se, cs) + | isCt0 w = + xs + | isCt1 w = + fuseConcatArg xs (ArgArrayLit [se], cs) +fuseConcatArg ((ArgArrayLit x_ses, x_cs) : xs) (ArgArrayLit y_ses, y_cs) = + (ArgArrayLit (x_ses ++ y_ses), x_cs <> y_cs) : xs +fuseConcatArg ((ArgReplicate x_ws x_se, x_cs) : xs) (ArgReplicate y_ws y_se, y_cs) + | x_se == y_se = + (ArgReplicate (x_ws ++ y_ws) x_se, x_cs <> y_cs) : xs +fuseConcatArg xs y = + y : xs + +simplifyConcat :: BinderOps lore => BottomUpRuleBasicOp lore +-- concat@1(transpose(x),transpose(y)) == transpose(concat@0(x,y)) +simplifyConcat (vtable, _) pat _ (Concat i x xs new_d) + | Just r <- arrayRank <$> ST.lookupType x vtable, + let perm = [i] ++ [0 .. i -1] ++ [i + 1 .. r -1], + Just (x', x_cs) <- transposedBy perm x, + Just (xs', xs_cs) <- unzip <$> mapM (transposedBy perm) xs = Simplify $ do + concat_rearrange <- + certifying (x_cs <> mconcat xs_cs) $ + letExp "concat_rearrange" $ BasicOp $ Concat 0 x' xs' new_d + letBind pat $ BasicOp $ Rearrange perm concat_rearrange + where + transposedBy perm1 v = + case ST.lookupExp v vtable of + Just (BasicOp (Rearrange perm2 v'), vcs) + | perm1 == perm2 -> Just (v', vcs) + _ -> Nothing + +-- Removing a concatenation that involves only a single array. This +-- may be produced as a result of other simplification rules. +simplifyConcat _ pat aux (Concat _ x [] _) = + Simplify $ + -- Still need a copy because Concat produces a fresh array. + auxing aux $ letBind pat $ BasicOp $ Copy x +-- concat xs (concat ys zs) == concat xs ys zs +simplifyConcat (vtable, _) pat (StmAux cs attrs _) (Concat i x xs new_d) + | x' /= x || concat xs' /= xs = + Simplify $ + certifying (cs <> x_cs <> mconcat xs_cs) $ + attributing attrs $ + letBind pat $ + BasicOp $ Concat i x' (zs ++ concat xs') new_d + where + (x' : zs, x_cs) = isConcat x + (xs', xs_cs) = unzip $ map isConcat xs + isConcat v = case ST.lookupBasicOp v vtable of + Just (Concat j y ys _, v_cs) | j == i -> (y : ys, v_cs) + _ -> ([v], mempty) + +-- Fusing arguments to the concat when possible. Only done when +-- concatenating along the outer dimension for now. +simplifyConcat (vtable, _) pat aux (Concat 0 x xs outer_w) + | -- We produce the to-be-concatenated arrays in reverse order, so + -- reverse them back. + y : ys <- + reverse $ + foldl' fuseConcatArg mempty $ + map (toConcatArg vtable) $ x : xs, + length xs /= length ys = + Simplify $ do + elem_type <- lookupType x + y' <- fromConcatArg elem_type y + ys' <- mapM (fromConcatArg elem_type) ys + auxing aux $ letBind pat $ BasicOp $ Concat 0 y' ys' outer_w +simplifyConcat _ _ _ _ = Skip + +ruleBasicOp :: BinderOps lore => TopDownRuleBasicOp lore +ruleBasicOp vtable pat aux op + | Just (op', cs) <- applySimpleRules defOf seType op = + Simplify $ certifying (cs <> stmAuxCerts aux) $ letBind pat $ BasicOp op' + where + defOf = (`ST.lookupExp` vtable) + seType (Var v) = ST.lookupType v vtable + seType (Constant v) = Just $ Prim $ primValueType v +ruleBasicOp vtable pat _ (Update src _ (Var v)) + | Just (BasicOp Scratch {}, _) <- ST.lookupExp v vtable = + Simplify $ letBind pat $ BasicOp $ SubExp $ Var src +-- If we are writing a single-element slice from some array, and the +-- element of that array can be computed as a PrimExp based on the +-- index, let's just write that instead. +ruleBasicOp vtable pat aux (Update src [DimSlice i n s] (Var v)) + | isCt1 n, + isCt1 s, + Just (ST.Indexed cs e) <- ST.index v [intConst Int64 0] vtable = + Simplify $ do + e' <- toSubExp "update_elem" e + auxing aux $ + certifying cs $ + letBind pat $ BasicOp $ Update src [DimFix i] e' +ruleBasicOp vtable pat _ (Update dest destis (Var v)) + | Just (e, _) <- ST.lookupExp v vtable, + arrayFrom e = + Simplify $ letBind pat $ BasicOp $ SubExp $ Var dest + where + arrayFrom (BasicOp (Copy copy_v)) + | Just (e', _) <- ST.lookupExp copy_v vtable = + arrayFrom e' + arrayFrom (BasicOp (Index src srcis)) = + src == dest && destis == srcis + arrayFrom (BasicOp (Replicate v_shape v_se)) + | Just (Replicate dest_shape dest_se, _) <- ST.lookupBasicOp dest vtable, + v_se == dest_se, + shapeDims v_shape `isSuffixOf` shapeDims dest_shape = + True + arrayFrom _ = + False +ruleBasicOp vtable pat _ (Update dest is se) + | Just dest_t <- ST.lookupType dest vtable, + isFullSlice (arrayShape dest_t) is = Simplify $ + case se of + Var v | not $ null $ sliceDims is -> do + v_reshaped <- + letExp (baseString v ++ "_reshaped") $ + BasicOp $ Reshape (map DimNew $ arrayDims dest_t) v + letBind pat $ BasicOp $ Copy v_reshaped + _ -> letBind pat $ BasicOp $ ArrayLit [se] $ rowType dest_t +ruleBasicOp vtable pat (StmAux cs1 attrs _) (Update dest1 is1 (Var v1)) + | Just (Update dest2 is2 se2, cs2) <- ST.lookupBasicOp v1 vtable, + Just (Copy v3, cs3) <- ST.lookupBasicOp dest2 vtable, + Just (Index v4 is4, cs4) <- ST.lookupBasicOp v3 vtable, + is4 == is1, + v4 == dest1 = + Simplify $ + certifying (cs1 <> cs2 <> cs3 <> cs4) $ do + is5 <- subExpSlice $ sliceSlice (primExpSlice is1) (primExpSlice is2) + attributing attrs $ letBind pat $ BasicOp $ Update dest1 is5 se2 +ruleBasicOp vtable pat _ (CmpOp (CmpEq t) se1 se2) + | Just m <- simplifyWith se1 se2 = Simplify m + | Just m <- simplifyWith se2 se1 = Simplify m + where + simplifyWith (Var v) x + | Just bnd <- ST.lookupStm v vtable, + If p tbranch fbranch _ <- stmExp bnd, + Just (y, z) <- + returns v (stmPattern bnd) tbranch fbranch, + not $ boundInBody tbranch `namesIntersect` freeIn y, + not $ boundInBody fbranch `namesIntersect` freeIn z = Just $ do + eq_x_y <- + letSubExp "eq_x_y" $ BasicOp $ CmpOp (CmpEq t) x y + eq_x_z <- + letSubExp "eq_x_z" $ BasicOp $ CmpOp (CmpEq t) x z + p_and_eq_x_y <- + letSubExp "p_and_eq_x_y" $ BasicOp $ BinOp LogAnd p eq_x_y + not_p <- + letSubExp "not_p" $ BasicOp $ UnOp Not p + not_p_and_eq_x_z <- + letSubExp "p_and_eq_x_y" $ BasicOp $ BinOp LogAnd not_p eq_x_z + letBind pat $ + BasicOp $ BinOp LogOr p_and_eq_x_y not_p_and_eq_x_z + simplifyWith _ _ = + Nothing + + returns v ifpat tbranch fbranch = + fmap snd $ + find ((== v) . patElemName . fst) $ + zip (patternValueElements ifpat) $ + zip (bodyResult tbranch) (bodyResult fbranch) +ruleBasicOp _ pat _ (Replicate (Shape []) se@Constant {}) = + Simplify $ letBind pat $ BasicOp $ SubExp se +ruleBasicOp _ pat _ (Replicate (Shape []) (Var v)) = Simplify $ do + v_t <- lookupType v + letBind pat $ + BasicOp $ + if primType v_t + then SubExp $ Var v + else Copy v +ruleBasicOp vtable pat _ (Replicate shape (Var v)) + | Just (BasicOp (Replicate shape2 se), cs) <- ST.lookupExp v vtable = + Simplify $ certifying cs $ letBind pat $ BasicOp $ Replicate (shape <> shape2) se +ruleBasicOp _ pat _ (ArrayLit (se : ses) _) + | all (== se) ses = + Simplify $ + let n = constant (fromIntegral (length ses) + 1 :: Int64) + in letBind pat $ BasicOp $ Replicate (Shape [n]) se +ruleBasicOp vtable pat aux (Index idd slice) + | Just inds <- sliceIndices slice, + Just (BasicOp (Reshape newshape idd2), idd_cs) <- ST.lookupExp idd vtable, + length newshape == length inds = + Simplify $ + case shapeCoercion newshape of + Just _ -> + certifying idd_cs $ + auxing aux $ + letBind pat $ BasicOp $ Index idd2 slice + Nothing -> do + -- Linearise indices and map to old index space. + oldshape <- arrayDims <$> lookupType idd2 + let new_inds = + reshapeIndex + (map pe64 oldshape) + (map pe64 $ newDims newshape) + (map pe64 inds) + new_inds' <- + mapM (toSubExp "new_index") new_inds + certifying idd_cs $ + auxing aux $ + letBind pat $ BasicOp $ Index idd2 $ map DimFix new_inds' +ruleBasicOp _ pat _ (BinOp (Pow t) e1 e2) + | e1 == intConst t 2 = + Simplify $ letBind pat $ BasicOp $ BinOp (Shl t) (intConst t 1) e2 +-- Handle identity permutation. +ruleBasicOp _ pat _ (Rearrange perm v) + | sort perm == perm = + Simplify $ letBind pat $ BasicOp $ SubExp $ Var v +ruleBasicOp vtable pat aux (Rearrange perm v) + | Just (BasicOp (Rearrange perm2 e), v_cs) <- ST.lookupExp v vtable = + -- Rearranging a rearranging: compose the permutations. + Simplify $ + certifying v_cs $ + auxing aux $ + letBind pat $ BasicOp $ Rearrange (perm `rearrangeCompose` perm2) e +ruleBasicOp vtable pat aux (Rearrange perm v) + | Just (BasicOp (Rotate offsets v2), v_cs) <- ST.lookupExp v vtable, + Just (BasicOp (Rearrange perm3 v3), v2_cs) <- ST.lookupExp v2 vtable = Simplify $ do + let offsets' = rearrangeShape (rearrangeInverse perm3) offsets + rearrange_rotate <- letExp "rearrange_rotate" $ BasicOp $ Rotate offsets' v3 + certifying (v_cs <> v2_cs) $ + auxing aux $ + letBind pat $ BasicOp $ Rearrange (perm `rearrangeCompose` perm3) rearrange_rotate + +-- Rearranging a replicate where the outer dimension is left untouched. +ruleBasicOp vtable pat aux (Rearrange perm v1) + | Just (BasicOp (Replicate dims (Var v2)), v1_cs) <- ST.lookupExp v1 vtable, + num_dims <- shapeRank dims, + (rep_perm, rest_perm) <- splitAt num_dims perm, + not $ null rest_perm, + rep_perm == [0 .. length rep_perm -1] = + Simplify $ + certifying v1_cs $ + auxing aux $ do + v <- + letSubExp "rearrange_replicate" $ + BasicOp $ Rearrange (map (subtract num_dims) rest_perm) v2 + letBind pat $ BasicOp $ Replicate dims v + +-- A zero-rotation is identity. +ruleBasicOp _ pat _ (Rotate offsets v) + | all isCt0 offsets = Simplify $ letBind pat $ BasicOp $ SubExp $ Var v +ruleBasicOp vtable pat aux (Rotate offsets v) + | Just (BasicOp (Rearrange perm v2), v_cs) <- ST.lookupExp v vtable, + Just (BasicOp (Rotate offsets2 v3), v2_cs) <- ST.lookupExp v2 vtable = Simplify $ do + let offsets2' = rearrangeShape (rearrangeInverse perm) offsets2 + addOffsets x y = letSubExp "summed_offset" $ BasicOp $ BinOp (Add Int64 OverflowWrap) x y + offsets' <- zipWithM addOffsets offsets offsets2' + rotate_rearrange <- + auxing aux $ letExp "rotate_rearrange" $ BasicOp $ Rearrange perm v3 + certifying (v_cs <> v2_cs) $ + letBind pat $ BasicOp $ Rotate offsets' rotate_rearrange + +-- Combining Rotates. +ruleBasicOp vtable pat aux (Rotate offsets1 v) + | Just (BasicOp (Rotate offsets2 v2), v_cs) <- ST.lookupExp v vtable = Simplify $ do + offsets <- zipWithM add offsets1 offsets2 + certifying v_cs $ + auxing aux $ + letBind pat $ BasicOp $ Rotate offsets v2 + where + add x y = letSubExp "offset" $ BasicOp $ BinOp (Add Int64 OverflowWrap) x y + +-- If we see an Update with a scalar where the value to be written is +-- the result of indexing some other array, then we convert it into an +-- Update with a slice of that array. This matters when the arrays +-- are far away (on the GPU, say), because it avoids a copy of the +-- scalar to and from the host. +ruleBasicOp vtable pat aux (Update arr_x slice_x (Var v)) + | Just _ <- sliceIndices slice_x, + Just (Index arr_y slice_y, cs_y) <- ST.lookupBasicOp v vtable, + ST.available arr_y vtable, + -- XXX: we should check for proper aliasing here instead. + arr_y /= arr_x, + Just (slice_x_bef, DimFix i, []) <- focusNth (length slice_x - 1) slice_x, + Just (slice_y_bef, DimFix j, []) <- focusNth (length slice_y - 1) slice_y = Simplify $ do + let slice_x' = slice_x_bef ++ [DimSlice i (intConst Int64 1) (intConst Int64 1)] + slice_y' = slice_y_bef ++ [DimSlice j (intConst Int64 1) (intConst Int64 1)] + v' <- letExp (baseString v ++ "_slice") $ BasicOp $ Index arr_y slice_y' + certifying cs_y $ + auxing aux $ + letBind pat $ BasicOp $ Update arr_x slice_x' $ Var v' + +-- Simplify away 0<=i when 'i' is from a loop of form 'for i < n'. +ruleBasicOp vtable pat aux (CmpOp CmpSle {} x y) + | Constant (IntValue (Int64Value 0)) <- x, + Var v <- y, + Just _ <- ST.lookupLoopVar v vtable = + Simplify $ auxing aux $ letBind pat $ BasicOp $ SubExp $ constant True +-- Simplify away i [TopDownRule lore] +topDownRules = + [ RuleBasicOp ruleBasicOp + ] + +bottomUpRules :: BinderOps lore => [BottomUpRule lore] +bottomUpRules = + [ RuleBasicOp simplifyConcat + ] + +-- | A set of simplification rules for 'BasicOp's. Includes rules +-- from "Futhark.Optimise.Simplify.Rules.Simple". +basicOpRules :: (BinderOps lore, Aliased lore) => RuleBook lore +basicOpRules = ruleBook topDownRules bottomUpRules <> loopRules diff --git a/src/Futhark/Optimise/Simplify/Rules/ClosedForm.hs b/src/Futhark/Optimise/Simplify/Rules/ClosedForm.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/Simplify/Rules/ClosedForm.hs @@ -0,0 +1,223 @@ +{-# LANGUAGE FlexibleContexts #-} + +-- | This module implements facilities for determining whether a +-- reduction or fold can be expressed in a closed form (i.e. not as a +-- SOAC). +-- +-- Right now, the module can detect only trivial cases. In the +-- future, we would like to make it more powerful, as well as possibly +-- also being able to analyse sequential loops. +module Futhark.Optimise.Simplify.Rules.ClosedForm + ( foldClosedForm, + loopClosedForm, + ) +where + +import Control.Monad +import qualified Data.Map.Strict as M +import Data.Maybe +import Futhark.Construct +import Futhark.IR +import Futhark.Optimise.Simplify.Rule +import Futhark.Optimise.Simplify.Rules.Simple (VarLookup) +import Futhark.Transform.Rename + +{- +Motivation: + + let {*[int,x_size_27] map_computed_shape_1286} = replicate(x_size_27, + all_equal_shape_1044) in + let {*[bool,x_size_27] map_size_checks_1292} = replicate(x_size_27, x_1291) in + let {bool all_equal_checked_1298, int all_equal_shape_1299} = + reduceT(fn {bool, int} (bool bacc_1293, int nacc_1294, bool belm_1295, + int nelm_1296) => + let {bool tuplit_elems_1297} = bacc_1293 && belm_1295 in + {tuplit_elems_1297, nelm_1296}, + {True, 0}, map_size_checks_1292, map_computed_shape_1286) +-} + +-- | @foldClosedForm look foldfun accargs arrargs@ determines whether +-- each of the results of @foldfun@ can be expressed in a closed form. +foldClosedForm :: + (ASTLore lore, BinderOps lore) => + VarLookup lore -> + Pattern lore -> + Lambda lore -> + [SubExp] -> + [VName] -> + RuleM lore () +foldClosedForm look pat lam accs arrs = do + inputsize <- arraysSize 0 <$> mapM lookupType arrs + + t <- case patternTypes pat of + [Prim t] -> return t + _ -> cannotSimplify + + closedBody <- + checkResults + (patternNames pat) + inputsize + mempty + Int64 + knownBnds + (map paramName (lambdaParams lam)) + (lambdaBody lam) + accs + isEmpty <- newVName "fold_input_is_empty" + letBindNames [isEmpty] $ + BasicOp $ CmpOp (CmpEq int64) inputsize (intConst Int64 0) + letBind pat + =<< ( If (Var isEmpty) + <$> resultBodyM accs + <*> renameBody closedBody + <*> pure (IfDec [primBodyType t] IfNormal) + ) + where + knownBnds = determineKnownBindings look lam accs arrs + +-- | @loopClosedForm pat respat merge bound bodys@ determines whether +-- the do-loop can be expressed in a closed form. +loopClosedForm :: + (ASTLore lore, BinderOps lore) => + Pattern lore -> + [(FParam lore, SubExp)] -> + Names -> + IntType -> + SubExp -> + Body lore -> + RuleM lore () +loopClosedForm pat merge i it bound body = do + t <- case patternTypes pat of + [Prim t] -> return t + _ -> cannotSimplify + + closedBody <- + checkResults + mergenames + bound + i + it + knownBnds + (map identName mergeidents) + body + mergeexp + isEmpty <- newVName "bound_is_zero" + letBindNames [isEmpty] $ + BasicOp $ CmpOp (CmpSlt it) bound (intConst it 0) + + letBind pat + =<< ( If (Var isEmpty) + <$> resultBodyM mergeexp + <*> renameBody closedBody + <*> pure (IfDec [primBodyType t] IfNormal) + ) + where + (mergepat, mergeexp) = unzip merge + mergeidents = map paramIdent mergepat + mergenames = map paramName mergepat + knownBnds = M.fromList $ zip mergenames mergeexp + +checkResults :: + BinderOps lore => + [VName] -> + SubExp -> + Names -> + IntType -> + M.Map VName SubExp -> + -- | Lambda-bound + [VName] -> + Body lore -> + [SubExp] -> + RuleM lore (Body lore) +checkResults pat size untouchable it knownBnds params body accs = do + ((), bnds) <- + collectStms $ + zipWithM_ checkResult (zip pat res) (zip accparams accs) + mkBodyM bnds $ map Var pat + where + bndMap = makeBindMap body + (accparams, _) = splitAt (length accs) params + res = bodyResult body + + nonFree = boundInBody body <> namesFromList params <> untouchable + + checkResult (p, Var v) (accparam, acc) + | Just (BasicOp (BinOp bop x y)) <- M.lookup v bndMap = do + -- One of x,y must be *this* accumulator, and the other must + -- be something that is free in the body. + let isThisAccum = (== Var accparam) + (this, el) <- liftMaybe $ + case ( (asFreeSubExp x, isThisAccum y), + (asFreeSubExp y, isThisAccum x) + ) of + ((Just free, True), _) -> Just (acc, free) + (_, (Just free, True)) -> Just (acc, free) + _ -> Nothing + + case bop of + LogAnd -> + letBindNames [p] $ BasicOp $ BinOp LogAnd this el + Add t w | Just properly_typed_size <- properIntSize t -> do + size' <- properly_typed_size + letBindNames [p] + =<< eBinOp + (Add t w) + (eSubExp this) + (pure $ BasicOp $ BinOp (Mul t w) el size') + FAdd t | Just properly_typed_size <- properFloatSize t -> do + size' <- properly_typed_size + letBindNames [p] + =<< eBinOp + (FAdd t) + (eSubExp this) + (pure $ BasicOp $ BinOp (FMul t) el size') + _ -> cannotSimplify -- Um... sorry. + checkResult _ _ = cannotSimplify + + asFreeSubExp :: SubExp -> Maybe SubExp + asFreeSubExp (Var v) + | v `nameIn` nonFree = M.lookup v knownBnds + asFreeSubExp se = Just se + + properIntSize Int64 = Just $ return size + properIntSize t = + Just $ + letSubExp "converted_size" $ + BasicOp $ ConvOp (SExt it t) size + + properFloatSize t = + Just $ + letSubExp "converted_size" $ + BasicOp $ ConvOp (SIToFP it t) size + +determineKnownBindings :: + VarLookup lore -> + Lambda lore -> + [SubExp] -> + [VName] -> + M.Map VName SubExp +determineKnownBindings look lam accs arrs = + accBnds <> arrBnds + where + (accparams, arrparams) = + splitAt (length accs) $ lambdaParams lam + accBnds = + M.fromList $ + zip (map paramName accparams) accs + arrBnds = + M.fromList $ + mapMaybe isReplicate $ + zip (map paramName arrparams) arrs + + isReplicate (p, v) + | Just (BasicOp (Replicate _ ve), cs) <- look v, + cs == mempty = + Just (p, ve) + isReplicate _ = Nothing + +makeBindMap :: Body lore -> M.Map VName (Exp lore) +makeBindMap = M.fromList . mapMaybe isSingletonStm . stmsToList . bodyStms + where + isSingletonStm (Let pat _ e) = case patternNames pat of + [v] -> Just (v, e) + _ -> Nothing diff --git a/src/Futhark/Optimise/Simplify/Rules/Index.hs b/src/Futhark/Optimise/Simplify/Rules/Index.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/Simplify/Rules/Index.hs @@ -0,0 +1,230 @@ +{-# OPTIONS_GHC -Wno-overlapping-patterns -Wno-incomplete-patterns -Wno-incomplete-uni-patterns -Wno-incomplete-record-updates #-} + +-- | Index simplification mechanics. +module Futhark.Optimise.Simplify.Rules.Index + ( IndexResult (..), + simplifyIndexing, + ) +where + +import Data.Maybe +import Futhark.Analysis.PrimExp.Convert +import qualified Futhark.Analysis.SymbolTable as ST +import Futhark.Construct +import Futhark.IR +import Futhark.Optimise.Simplify.Rules.Simple +import Futhark.Util + +isCt1 :: SubExp -> Bool +isCt1 (Constant v) = oneIsh v +isCt1 _ = False + +isCt0 :: SubExp -> Bool +isCt0 (Constant v) = zeroIsh v +isCt0 _ = False + +-- | Some index expressions can be simplified to 'SubExp's, while +-- others produce another index expression (which may be further +-- simplifiable). +data IndexResult + = IndexResult Certificates VName (Slice SubExp) + | SubExpResult Certificates SubExp + +-- | Try to simplify an index operation. +simplifyIndexing :: + MonadBinder m => + ST.SymbolTable (Lore m) -> + TypeLookup -> + VName -> + Slice SubExp -> + Bool -> + Maybe (m IndexResult) +simplifyIndexing vtable seType idd inds consuming = + case defOf idd of + _ + | Just t <- seType (Var idd), + inds == fullSlice t [] -> + Just $ pure $ SubExpResult mempty $ Var idd + | Just inds' <- sliceIndices inds, + Just (ST.Indexed cs e) <- ST.index idd inds' vtable, + worthInlining e, + all (`ST.elem` vtable) (unCertificates cs) -> + Just $ SubExpResult cs <$> toSubExp "index_primexp" e + | Just inds' <- sliceIndices inds, + Just (ST.IndexedArray cs arr inds'') <- ST.index idd inds' vtable, + all (worthInlining . untyped) inds'', + all (`ST.elem` vtable) (unCertificates cs) -> + Just $ + IndexResult cs arr . map DimFix + <$> mapM (toSubExp "index_primexp") inds'' + Nothing -> Nothing + Just (SubExp (Var v), cs) -> Just $ pure $ IndexResult cs v inds + Just (Iota _ x s to_it, cs) + | [DimFix ii] <- inds, + Just (Prim (IntType from_it)) <- seType ii -> + Just $ + let mul = BinOpExp $ Mul to_it OverflowWrap + add = BinOpExp $ Add to_it OverflowWrap + in fmap (SubExpResult cs) $ + toSubExp "index_iota" $ + ( sExt to_it (primExpFromSubExp (IntType from_it) ii) + `mul` primExpFromSubExp (IntType to_it) s + ) + `add` primExpFromSubExp (IntType to_it) x + | [DimSlice i_offset i_n i_stride] <- inds -> + Just $ do + i_offset' <- asIntS to_it i_offset + i_stride' <- asIntS to_it i_stride + let mul = BinOpExp $ Mul to_it OverflowWrap + add = BinOpExp $ Add to_it OverflowWrap + i_offset'' <- + toSubExp "iota_offset" $ + ( primExpFromSubExp (IntType to_it) x + `mul` primExpFromSubExp (IntType to_it) s + ) + `add` primExpFromSubExp (IntType to_it) i_offset' + i_stride'' <- + letSubExp "iota_offset" $ + BasicOp $ BinOp (Mul Int64 OverflowWrap) s i_stride' + fmap (SubExpResult cs) $ + letSubExp "slice_iota" $ + BasicOp $ Iota i_n i_offset'' i_stride'' to_it + + -- A rotate cannot be simplified away if we are slicing a rotated dimension. + Just (Rotate offsets a, cs) + | not $ or $ zipWith rotateAndSlice offsets inds -> Just $ do + dims <- arrayDims <$> lookupType a + let adjustI i o d = do + i_p_o <- letSubExp "i_p_o" $ BasicOp $ BinOp (Add Int64 OverflowWrap) i o + letSubExp "rot_i" (BasicOp $ BinOp (SMod Int64 Unsafe) i_p_o d) + adjust (DimFix i, o, d) = + DimFix <$> adjustI i o d + adjust (DimSlice i n s, o, d) = + DimSlice <$> adjustI i o d <*> pure n <*> pure s + IndexResult cs a <$> mapM adjust (zip3 inds offsets dims) + where + rotateAndSlice r DimSlice {} = not $ isCt0 r + rotateAndSlice _ _ = False + Just (Index aa ais, cs) -> + Just $ + IndexResult cs aa + <$> subExpSlice (sliceSlice (primExpSlice ais) (primExpSlice inds)) + Just (Replicate (Shape [_]) (Var vv), cs) + | [DimFix {}] <- inds, not consuming -> Just $ pure $ SubExpResult cs $ Var vv + | DimFix {} : is' <- inds, not consuming -> Just $ pure $ IndexResult cs vv is' + Just (Replicate (Shape [_]) val@(Constant _), cs) + | [DimFix {}] <- inds, not consuming -> Just $ pure $ SubExpResult cs val + Just (Replicate (Shape ds) v, cs) + | (ds_inds, rest_inds) <- splitAt (length ds) inds, + (ds', ds_inds') <- unzip $ mapMaybe index ds_inds, + ds' /= ds -> + Just $ do + arr <- letExp "smaller_replicate" $ BasicOp $ Replicate (Shape ds') v + return $ IndexResult cs arr $ ds_inds' ++ rest_inds + where + index DimFix {} = Nothing + index (DimSlice _ n s) = Just (n, DimSlice (constant (0 :: Int64)) n s) + Just (Rearrange perm src, cs) + | rearrangeReach perm <= length (takeWhile isIndex inds) -> + let inds' = rearrangeShape (rearrangeInverse perm) inds + in Just $ pure $ IndexResult cs src inds' + where + isIndex DimFix {} = True + isIndex _ = False + Just (Copy src, cs) + | Just dims <- arrayDims <$> seType (Var src), + length inds == length dims, + -- It is generally not safe to simplify a slice of a copy, + -- because the result may be used in an in-place update of the + -- original. + Just _ <- mapM dimFix inds, + not consuming, + ST.available src vtable -> + Just $ pure $ IndexResult cs src inds + Just (Reshape newshape src, cs) + | Just newdims <- shapeCoercion newshape, + Just olddims <- arrayDims <$> seType (Var src), + changed_dims <- zipWith (/=) newdims olddims, + not $ or $ drop (length inds) changed_dims -> + Just $ pure $ IndexResult cs src inds + | Just newdims <- shapeCoercion newshape, + Just olddims <- arrayDims <$> seType (Var src), + length newshape == length inds, + length olddims == length newdims -> + Just $ pure $ IndexResult cs src inds + Just (Reshape [_] v2, cs) + | Just [_] <- arrayDims <$> seType (Var v2) -> + Just $ pure $ IndexResult cs v2 inds + Just (Concat d x xs _, cs) + | -- HACK: simplifying the indexing of an N-array concatenation + -- is going to produce an N-deep if expression, which is bad + -- when N is large. To try to avoid that, we use the + -- heuristic not to simplify as long as any of the operands + -- are themselves Concats. The hope it that this will give + -- simplification some time to cut down the concatenation to + -- something smaller, before we start inlining. + not $ any isConcat $ x : xs, + Just (ibef, DimFix i, iaft) <- focusNth d inds, + Just (Prim res_t) <- + (`setArrayDims` sliceDims inds) + <$> ST.lookupType x vtable -> Just $ do + x_len <- arraySize d <$> lookupType x + xs_lens <- mapM (fmap (arraySize d) . lookupType) xs + + let add n m = do + added <- letSubExp "index_concat_add" $ BasicOp $ BinOp (Add Int64 OverflowWrap) n m + return (added, n) + (_, starts) <- mapAccumLM add x_len xs_lens + let xs_and_starts = reverse $ zip xs starts + + let mkBranch [] = + letSubExp "index_concat" $ BasicOp $ Index x $ ibef ++ DimFix i : iaft + mkBranch ((x', start) : xs_and_starts') = do + cmp <- letSubExp "index_concat_cmp" $ BasicOp $ CmpOp (CmpSle Int64) start i + (thisres, thisbnds) <- collectStms $ do + i' <- letSubExp "index_concat_i" $ BasicOp $ BinOp (Sub Int64 OverflowWrap) i start + letSubExp "index_concat" $ BasicOp $ Index x' $ ibef ++ DimFix i' : iaft + thisbody <- mkBodyM thisbnds [thisres] + (altres, altbnds) <- collectStms $ mkBranch xs_and_starts' + altbody <- mkBodyM altbnds [altres] + letSubExp "index_concat_branch" $ + If cmp thisbody altbody $ + IfDec [primBodyType res_t] IfNormal + SubExpResult cs <$> mkBranch xs_and_starts + Just (ArrayLit ses _, cs) + | DimFix (Constant (IntValue (Int64Value i))) : inds' <- inds, + Just se <- maybeNth i ses -> + case inds' of + [] -> Just $ pure $ SubExpResult cs se + _ | Var v2 <- se -> Just $ pure $ IndexResult cs v2 inds' + _ -> Nothing + -- Indexing single-element arrays. We know the index must be 0. + _ + | Just t <- seType $ Var idd, + isCt1 $ arraySize 0 t, + DimFix i : inds' <- inds, + not $ isCt0 i -> + Just . pure . IndexResult mempty idd $ + DimFix (constant (0 :: Int64)) : inds' + _ -> Nothing + where + defOf v = do + (BasicOp op, def_cs) <- ST.lookupExp v vtable + return (op, def_cs) + worthInlining e + | primExpSizeAtLeast 20 e = False -- totally ad-hoc. + | otherwise = worthInlining' e + worthInlining' (BinOpExp Pow {} _ _) = False + worthInlining' (BinOpExp FPow {} _ _) = False + worthInlining' (BinOpExp _ x y) = worthInlining' x && worthInlining' y + worthInlining' (CmpOpExp _ x y) = worthInlining' x && worthInlining' y + worthInlining' (ConvOpExp _ x) = worthInlining' x + worthInlining' (UnOpExp _ x) = worthInlining' x + worthInlining' FunExp {} = False + worthInlining' _ = True + + isConcat v + | Just (Concat {}, _) <- defOf v = + True + | otherwise = + False diff --git a/src/Futhark/Optimise/Simplify/Rules/Loop.hs b/src/Futhark/Optimise/Simplify/Rules/Loop.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/Simplify/Rules/Loop.hs @@ -0,0 +1,378 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | Loop simplification rules. +module Futhark.Optimise.Simplify.Rules.Loop (loopRules) where + +import Control.Monad +import Data.List (partition) +import Data.Maybe +import Futhark.Analysis.DataDependencies +import Futhark.Analysis.PrimExp.Convert +import qualified Futhark.Analysis.SymbolTable as ST +import qualified Futhark.Analysis.UsageTable as UT +import Futhark.Construct +import Futhark.IR +import Futhark.IR.Prop.Aliases +import Futhark.Optimise.Simplify.Rule +import Futhark.Optimise.Simplify.Rules.ClosedForm +import Futhark.Optimise.Simplify.Rules.Index +import Futhark.Transform.Rename + +-- This next one is tricky - it's easy enough to determine that some +-- loop result is not used after the loop, but here, we must also make +-- sure that it does not affect any other values. +-- +-- I do not claim that the current implementation of this rule is +-- perfect, but it should suffice for many cases, and should never +-- generate wrong code. +removeRedundantMergeVariables :: BinderOps lore => BottomUpRuleDoLoop lore +removeRedundantMergeVariables (_, used) pat aux (ctx, val, form, body) + | not $ all (usedAfterLoop . fst) val, + null ctx -- FIXME: things get tricky if we can remove all vals + -- but some ctxs are still used. We take the easy way + -- out for now. + = + let (ctx_es, val_es) = splitAt (length ctx) $ bodyResult body + necessaryForReturned = + findNecessaryForReturned + usedAfterLoopOrInForm + (zip (map fst $ ctx ++ val) $ ctx_es ++ val_es) + (dataDependencies body) + + resIsNecessary ((v, _), _) = + usedAfterLoop v + || paramName v `nameIn` necessaryForReturned + || referencedInPat v + || referencedInForm v + + (keep_ctx, discard_ctx) = + partition resIsNecessary $ zip ctx ctx_es + (keep_valpart, discard_valpart) = + partition (resIsNecessary . snd) $ + zip (patternValueElements pat) $ zip val val_es + + (keep_valpatelems, keep_val) = unzip keep_valpart + (_discard_valpatelems, discard_val) = unzip discard_valpart + (ctx', ctx_es') = unzip keep_ctx + (val', val_es') = unzip keep_val + + body' = body {bodyResult = ctx_es' ++ val_es'} + free_in_keeps = freeIn keep_valpatelems + + stillUsedContext pat_elem = + patElemName pat_elem + `nameIn` ( free_in_keeps + <> freeIn (filter (/= pat_elem) $ patternContextElements pat) + ) + + pat' = + pat + { patternValueElements = keep_valpatelems, + patternContextElements = + filter stillUsedContext $ patternContextElements pat + } + in if ctx' ++ val' == ctx ++ val + then Skip + else Simplify $ do + -- We can't just remove the bindings in 'discard', since the loop + -- body may still use their names in (now-dead) expressions. + -- Hence, we add them inside the loop, fully aware that dead-code + -- removal will eventually get rid of them. Some care is + -- necessary to handle unique bindings. + body'' <- insertStmsM $ do + mapM_ (uncurry letBindNames) $ dummyStms discard_ctx + mapM_ (uncurry letBindNames) $ dummyStms discard_val + return body' + auxing aux $ letBind pat' $ DoLoop ctx' val' form body'' + where + pat_used = map (`UT.isUsedDirectly` used) $ patternValueNames pat + used_vals = map fst $ filter snd $ zip (map (paramName . fst) val) pat_used + usedAfterLoop = flip elem used_vals . paramName + usedAfterLoopOrInForm p = + usedAfterLoop p || paramName p `nameIn` freeIn form + patAnnotNames = freeIn $ map fst $ ctx ++ val + referencedInPat = (`nameIn` patAnnotNames) . paramName + referencedInForm = (`nameIn` freeIn form) . paramName + + dummyStms = map dummyStm + dummyStm ((p, e), _) + | unique (paramDeclType p), + Var v <- e = + ([paramName p], BasicOp $ Copy v) + | otherwise = ([paramName p], BasicOp $ SubExp e) +removeRedundantMergeVariables _ _ _ _ = + Skip + +-- We may change the type of the loop if we hoist out a shape +-- annotation, in which case we also need to tweak the bound pattern. +hoistLoopInvariantMergeVariables :: BinderOps lore => TopDownRuleDoLoop lore +hoistLoopInvariantMergeVariables vtable pat aux (ctx, val, form, loopbody) = + -- Figure out which of the elements of loopresult are + -- loop-invariant, and hoist them out. + case foldr checkInvariance ([], explpat, [], []) $ + zip3 (patternNames pat) merge res of + ([], _, _, _) -> + -- Nothing is invariant. + Skip + (invariant, explpat', merge', res') -> Simplify $ do + -- We have moved something invariant out of the loop. + let loopbody' = loopbody {bodyResult = res'} + invariantShape :: (a, VName) -> Bool + invariantShape (_, shapemerge) = + shapemerge + `elem` map (paramName . fst) merge' + (implpat', implinvariant) = partition invariantShape implpat + implinvariant' = [(patElemIdent p, Var v) | (p, v) <- implinvariant] + implpat'' = map fst implpat' + explpat'' = map fst explpat' + (ctx', val') = splitAt (length implpat') merge' + forM_ (invariant ++ implinvariant') $ \(v1, v2) -> + letBindNames [identName v1] $ BasicOp $ SubExp v2 + auxing aux $ + letBind (Pattern implpat'' explpat'') $ + DoLoop ctx' val' form loopbody' + where + merge = ctx ++ val + res = bodyResult loopbody + + implpat = + zip (patternContextElements pat) $ + map (paramName . fst) ctx + explpat = + zip (patternValueElements pat) $ + map (paramName . fst) val + + namesOfMergeParams = namesFromList $ map (paramName . fst) $ ctx ++ val + + removeFromResult (mergeParam, mergeInit) explpat' = + case partition ((== paramName mergeParam) . snd) explpat' of + ([(patelem, _)], rest) -> + (Just (patElemIdent patelem, mergeInit), rest) + (_, _) -> + (Nothing, explpat') + + checkInvariance + (pat_name, (mergeParam, mergeInit), resExp) + (invariant, explpat', merge', resExps) + | not (unique (paramDeclType mergeParam)) + || arrayRank (paramDeclType mergeParam) == 1, + isInvariant, + -- Also do not remove the condition in a while-loop. + not $ paramName mergeParam `nameIn` freeIn form = + let (bnd, explpat'') = + removeFromResult (mergeParam, mergeInit) explpat' + in ( maybe id (:) bnd $ (paramIdent mergeParam, mergeInit) : invariant, + explpat'', + merge', + resExps + ) + where + -- A non-unique merge variable is invariant if one of the + -- following is true: + -- + -- (0) The result is a variable of the same name as the + -- parameter, where all existential parameters are already + -- known to be invariant + isInvariant + | Var v2 <- resExp, + paramName mergeParam == v2 = + allExistentialInvariant + (namesFromList $ map (identName . fst) invariant) + mergeParam + -- (1) The result is identical to the initial parameter value. + | mergeInit == resExp = True + -- (2) The initial parameter value is equal to an outer + -- loop parameter 'P', where the initial value of 'P' is + -- equal to 'resExp', AND 'resExp' ultimately becomes the + -- new value of 'P'. XXX: it's a bit clumsy that this + -- only works for one level of nesting, and I think it + -- would not be too hard to generalise. + | Var init_v <- mergeInit, + Just (p_init, p_res) <- ST.lookupLoopParam init_v vtable, + p_init == resExp, + p_res == Var pat_name = + True + | otherwise = False + checkInvariance + (_pat_name, (mergeParam, mergeInit), resExp) + (invariant, explpat', merge', resExps) = + (invariant, explpat', (mergeParam, mergeInit) : merge', resExp : resExps) + + allExistentialInvariant namesOfInvariant mergeParam = + all (invariantOrNotMergeParam namesOfInvariant) $ + namesToList $ + freeIn mergeParam `namesSubtract` oneName (paramName mergeParam) + invariantOrNotMergeParam namesOfInvariant name = + not (name `nameIn` namesOfMergeParams) + || name `nameIn` namesOfInvariant + +simplifyClosedFormLoop :: BinderOps lore => TopDownRuleDoLoop lore +simplifyClosedFormLoop _ pat _ ([], val, ForLoop i it bound [], body) = + Simplify $ loopClosedForm pat val (oneName i) it bound body +simplifyClosedFormLoop _ _ _ _ = Skip + +simplifyLoopVariables :: (BinderOps lore, Aliased lore) => TopDownRuleDoLoop lore +simplifyLoopVariables vtable pat aux (ctx, val, form@(ForLoop i it num_iters loop_vars), body) + | simplifiable <- map checkIfSimplifiable loop_vars, + not $ all isNothing simplifiable = Simplify $ do + -- Check if the simplifications throw away more information than + -- we are comfortable with at this stage. + (maybe_loop_vars, body_prefix_stms) <- + localScope (scopeOf form) $ + unzip <$> zipWithM onLoopVar loop_vars simplifiable + if maybe_loop_vars == map Just loop_vars + then cannotSimplify + else do + body' <- insertStmsM $ do + addStms $ mconcat body_prefix_stms + resultBodyM =<< bodyBind body + auxing aux $ + letBind pat $ + DoLoop + ctx + val + (ForLoop i it num_iters $ catMaybes maybe_loop_vars) + body' + where + seType (Var v) + | v == i = Just $ Prim $ IntType it + | otherwise = ST.lookupType v vtable + seType (Constant v) = Just $ Prim $ primValueType v + consumed_in_body = consumedInBody body + + vtable' = ST.fromScope (scopeOf form) <> vtable + + checkIfSimplifiable (p, arr) = + simplifyIndexing + vtable' + seType + arr + (DimFix (Var i) : fullSlice (paramType p) []) + $ paramName p `nameIn` consumed_in_body + + -- We only want this simplification if the result does not refer + -- to 'i' at all, or does not contain accesses. + onLoopVar (p, arr) Nothing = + return (Just (p, arr), mempty) + onLoopVar (p, arr) (Just m) = do + (x, x_stms) <- collectStms m + case x of + IndexResult cs arr' slice + | not $ any ((i `nameIn`) . freeIn) x_stms, + DimFix (Var j) : slice' <- slice, + j == i, + not $ i `nameIn` freeIn slice -> do + addStms x_stms + w <- arraySize 0 <$> lookupType arr' + for_in_partial <- + certifying cs $ + letExp "for_in_partial" $ + BasicOp $ + Index arr' $ + DimSlice (intConst Int64 0) w (intConst Int64 1) : slice' + return (Just (p, for_in_partial), mempty) + SubExpResult cs se + | all (notIndex . stmExp) x_stms -> do + x_stms' <- collectStms_ $ + certifying cs $ do + addStms x_stms + letBindNames [paramName p] $ BasicOp $ SubExp se + return (Nothing, x_stms') + _ -> return (Just (p, arr), mempty) + + notIndex (BasicOp Index {}) = False + notIndex _ = True +simplifyLoopVariables _ _ _ _ = Skip + +-- If a for-loop with no loop variables has a counter of type Int64, +-- and the bound is just a constant or sign-extended integer of +-- smaller type, then change the loop to iterate over the smaller type +-- instead. We then move the sign extension inside the loop instead. +-- This addresses loops of the form @for i in x.. TopDownRuleDoLoop lore +narrowLoopType vtable pat aux (ctx, val, ForLoop i Int64 n [], body) + | Just (n', it', cs) <- smallerType = + Simplify $ do + i' <- newVName $ baseString i + let form' = ForLoop i' it' n' [] + body' <- insertStmsM $ + inScopeOf form' $ do + letBindNames [i] $ BasicOp $ ConvOp (SExt it' Int64) (Var i') + pure body + auxing aux $ + certifying cs $ + letBind pat $ DoLoop ctx val form' body' + where + smallerType + | Var n' <- n, + Just (ConvOp (SExt it' _) n'', cs) <- ST.lookupBasicOp n' vtable = + Just (n'', it', cs) + | Constant (IntValue (Int64Value n')) <- n, + toInteger n' <= toInteger (maxBound :: Int32) = + Just (intConst Int32 (toInteger n'), Int32, mempty) + | otherwise = + Nothing +narrowLoopType _ _ _ _ = Skip + +unroll :: + BinderOps lore => + Integer -> + [(FParam lore, SubExp)] -> + (VName, IntType, Integer) -> + [(LParam lore, VName)] -> + Body lore -> + RuleM lore [SubExp] +unroll n merge (iv, it, i) loop_vars body + | i >= n = + return $ map snd merge + | otherwise = do + iter_body <- insertStmsM $ do + forM_ merge $ \(mergevar, mergeinit) -> + letBindNames [paramName mergevar] $ BasicOp $ SubExp mergeinit + + letBindNames [iv] $ BasicOp $ SubExp $ intConst it i + + forM_ loop_vars $ \(p, arr) -> + letBindNames [paramName p] $ + BasicOp $ + Index arr $ + DimFix (intConst Int64 i) : fullSlice (paramType p) [] + + -- Some of the sizes in the types here might be temporarily wrong + -- until copy propagation fixes it up. + pure body + + iter_body' <- renameBody iter_body + addStms $ bodyStms iter_body' + + let merge' = zip (map fst merge) $ bodyResult iter_body' + unroll n merge' (iv, it, i + 1) loop_vars body + +simplifyKnownIterationLoop :: BinderOps lore => TopDownRuleDoLoop lore +simplifyKnownIterationLoop _ pat aux (ctx, val, ForLoop i it (Constant iters) loop_vars, body) + | IntValue n <- iters, + zeroIshInt n || oneIshInt n || "unroll" `inAttrs` stmAuxAttrs aux = Simplify $ do + res <- unroll (valueIntegral n) (ctx ++ val) (i, it, 0) loop_vars body + forM_ (zip (patternNames pat) res) $ \(v, se) -> + letBindNames [v] $ BasicOp $ SubExp se +simplifyKnownIterationLoop _ _ _ _ = + Skip + +topDownRules :: (BinderOps lore, Aliased lore) => [TopDownRule lore] +topDownRules = + [ RuleDoLoop hoistLoopInvariantMergeVariables, + RuleDoLoop simplifyClosedFormLoop, + RuleDoLoop simplifyKnownIterationLoop, + RuleDoLoop simplifyLoopVariables, + RuleDoLoop narrowLoopType + ] + +bottomUpRules :: BinderOps lore => [BottomUpRule lore] +bottomUpRules = + [ RuleDoLoop removeRedundantMergeVariables + ] + +-- | Standard loop simplification rules. +loopRules :: (BinderOps lore, Aliased lore) => RuleBook lore +loopRules = ruleBook topDownRules bottomUpRules diff --git a/src/Futhark/Optimise/Simplify/Rules/Simple.hs b/src/Futhark/Optimise/Simplify/Rules/Simple.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/Simplify/Rules/Simple.hs @@ -0,0 +1,342 @@ +{-# LANGUAGE TupleSections #-} + +-- | Particularly simple simplification rules. +module Futhark.Optimise.Simplify.Rules.Simple + ( TypeLookup, + VarLookup, + applySimpleRules, + ) +where + +import Control.Monad +import Data.List (isSuffixOf) +import Futhark.Analysis.PrimExp.Convert +import Futhark.IR + +-- | A function that, given a variable name, returns its definition. +type VarLookup lore = VName -> Maybe (Exp lore, Certificates) + +-- | A function that, given a subexpression, returns its type. +type TypeLookup = SubExp -> Maybe Type + +-- | A simple rule is a top-down rule that can be expressed as a pure +-- function. +type SimpleRule lore = VarLookup lore -> TypeLookup -> BasicOp -> Maybe (BasicOp, Certificates) + +isCt1 :: SubExp -> Bool +isCt1 (Constant v) = oneIsh v +isCt1 _ = False + +isCt0 :: SubExp -> Bool +isCt0 (Constant v) = zeroIsh v +isCt0 _ = False + +simplifyCmpOp :: SimpleRule lore +simplifyCmpOp _ _ (CmpOp cmp e1 e2) + | e1 == e2 = constRes $ + BoolValue $ + case cmp of + CmpEq {} -> True + CmpSlt {} -> False + CmpUlt {} -> False + CmpSle {} -> True + CmpUle {} -> True + FCmpLt {} -> False + FCmpLe {} -> True + CmpLlt -> False + CmpLle -> True +simplifyCmpOp _ _ (CmpOp cmp (Constant v1) (Constant v2)) = + constRes . BoolValue =<< doCmpOp cmp v1 v2 +simplifyCmpOp look _ (CmpOp CmpEq {} (Constant (IntValue x)) (Var v)) + | Just (BasicOp (ConvOp BToI {} b), cs) <- look v = + case valueIntegral x :: Int of + 1 -> Just (SubExp b, cs) + 0 -> Just (UnOp Not b, cs) + _ -> Just (SubExp (Constant (BoolValue False)), cs) +simplifyCmpOp _ _ _ = Nothing + +simplifyBinOp :: SimpleRule lore +simplifyBinOp _ _ (BinOp op (Constant v1) (Constant v2)) + | Just res <- doBinOp op v1 v2 = + constRes res +simplifyBinOp look _ (BinOp Add {} e1 e2) + | isCt0 e1 = subExpRes e2 + | isCt0 e2 = subExpRes e1 + -- x+(y-x) => y + | Var v2 <- e2, + Just (BasicOp (BinOp Sub {} e2_a e2_b), cs) <- look v2, + e2_b == e1 = + Just (SubExp e2_a, cs) +simplifyBinOp _ _ (BinOp FAdd {} e1 e2) + | isCt0 e1 = subExpRes e2 + | isCt0 e2 = subExpRes e1 +simplifyBinOp look _ (BinOp sub@(Sub t _) e1 e2) + | isCt0 e2 = subExpRes e1 + -- Cases for simplifying (a+b)-b and permutations. + + -- (e1_a+e1_b)-e1_a == e1_b + | Var v1 <- e1, + Just (BasicOp (BinOp Add {} e1_a e1_b), cs) <- look v1, + e1_a == e2 = + Just (SubExp e1_b, cs) + -- (e1_a+e1_b)-e1_b == e1_a + | Var v1 <- e1, + Just (BasicOp (BinOp Add {} e1_a e1_b), cs) <- look v1, + e1_b == e2 = + Just (SubExp e1_a, cs) + -- e2_a-(e2_a+e2_b) == 0-e2_b + | Var v2 <- e2, + Just (BasicOp (BinOp Add {} e2_a e2_b), cs) <- look v2, + e2_a == e1 = + Just (BinOp sub (intConst t 0) e2_b, cs) + -- e2_b-(e2_a+e2_b) == 0-e2_a + | Var v2 <- e2, + Just (BasicOp (BinOp Add {} e2_a e2_b), cs) <- look v2, + e2_b == e1 = + Just (BinOp sub (intConst t 0) e2_a, cs) +simplifyBinOp _ _ (BinOp FSub {} e1 e2) + | isCt0 e2 = subExpRes e1 +simplifyBinOp _ _ (BinOp Mul {} e1 e2) + | isCt0 e1 = subExpRes e1 + | isCt0 e2 = subExpRes e2 + | isCt1 e1 = subExpRes e2 + | isCt1 e2 = subExpRes e1 +simplifyBinOp _ _ (BinOp FMul {} e1 e2) + | isCt0 e1 = subExpRes e1 + | isCt0 e2 = subExpRes e2 + | isCt1 e1 = subExpRes e2 + | isCt1 e2 = subExpRes e1 +simplifyBinOp look _ (BinOp (SMod t _) e1 e2) + | isCt1 e2 = constRes $ IntValue $ intValue t (0 :: Int) + | e1 == e2 = constRes $ IntValue $ intValue t (0 :: Int) + | Var v1 <- e1, + Just (BasicOp (BinOp SMod {} _ e4), v1_cs) <- look v1, + e4 == e2 = + Just (SubExp e1, v1_cs) +simplifyBinOp _ _ (BinOp SDiv {} e1 e2) + | isCt0 e1 = subExpRes e1 + | isCt1 e2 = subExpRes e1 + | isCt0 e2 = Nothing +simplifyBinOp _ _ (BinOp SDivUp {} e1 e2) + | isCt0 e1 = subExpRes e1 + | isCt1 e2 = subExpRes e1 + | isCt0 e2 = Nothing +simplifyBinOp _ _ (BinOp FDiv {} e1 e2) + | isCt0 e1 = subExpRes e1 + | isCt1 e2 = subExpRes e1 + | isCt0 e2 = Nothing +simplifyBinOp _ _ (BinOp (SRem t _) e1 e2) + | isCt1 e2 = constRes $ IntValue $ intValue t (0 :: Int) + | e1 == e2 = constRes $ IntValue $ intValue t (1 :: Int) +simplifyBinOp _ _ (BinOp SQuot {} e1 e2) + | isCt1 e2 = subExpRes e1 + | isCt0 e2 = Nothing +simplifyBinOp _ _ (BinOp (FPow t) e1 e2) + | isCt0 e2 = subExpRes $ floatConst t 1 + | isCt0 e1 || isCt1 e1 || isCt1 e2 = subExpRes e1 +simplifyBinOp _ _ (BinOp (Shl t) e1 e2) + | isCt0 e2 = subExpRes e1 + | isCt0 e1 = subExpRes $ intConst t 0 +simplifyBinOp _ _ (BinOp AShr {} e1 e2) + | isCt0 e2 = subExpRes e1 +simplifyBinOp _ _ (BinOp (And t) e1 e2) + | isCt0 e1 = subExpRes $ intConst t 0 + | isCt0 e2 = subExpRes $ intConst t 0 + | e1 == e2 = subExpRes e1 +simplifyBinOp _ _ (BinOp Or {} e1 e2) + | isCt0 e1 = subExpRes e2 + | isCt0 e2 = subExpRes e1 + | e1 == e2 = subExpRes e1 +simplifyBinOp _ _ (BinOp (Xor t) e1 e2) + | isCt0 e1 = subExpRes e2 + | isCt0 e2 = subExpRes e1 + | e1 == e2 = subExpRes $ intConst t 0 +simplifyBinOp defOf _ (BinOp LogAnd e1 e2) + | isCt0 e1 = constRes $ BoolValue False + | isCt0 e2 = constRes $ BoolValue False + | isCt1 e1 = subExpRes e2 + | isCt1 e2 = subExpRes e1 + | Var v <- e1, + Just (BasicOp (UnOp Not e1'), v_cs) <- defOf v, + e1' == e2 = + Just (SubExp $ Constant $ BoolValue False, v_cs) + | Var v <- e2, + Just (BasicOp (UnOp Not e2'), v_cs) <- defOf v, + e2' == e1 = + Just (SubExp $ Constant $ BoolValue False, v_cs) +simplifyBinOp defOf _ (BinOp LogOr e1 e2) + | isCt0 e1 = subExpRes e2 + | isCt0 e2 = subExpRes e1 + | isCt1 e1 = constRes $ BoolValue True + | isCt1 e2 = constRes $ BoolValue True + | Var v <- e1, + Just (BasicOp (UnOp Not e1'), v_cs) <- defOf v, + e1' == e2 = + Just (SubExp $ Constant $ BoolValue True, v_cs) + | Var v <- e2, + Just (BasicOp (UnOp Not e2'), v_cs) <- defOf v, + e2' == e1 = + Just (SubExp $ Constant $ BoolValue True, v_cs) +simplifyBinOp defOf _ (BinOp (SMax it) e1 e2) + | e1 == e2 = + subExpRes e1 + | Var v1 <- e1, + Just (BasicOp (BinOp (SMax _) e1_1 e1_2), v1_cs) <- defOf v1, + e1_1 == e2 = + Just (BinOp (SMax it) e1_2 e2, v1_cs) + | Var v1 <- e1, + Just (BasicOp (BinOp (SMax _) e1_1 e1_2), v1_cs) <- defOf v1, + e1_2 == e2 = + Just (BinOp (SMax it) e1_1 e2, v1_cs) + | Var v2 <- e2, + Just (BasicOp (BinOp (SMax _) e2_1 e2_2), v2_cs) <- defOf v2, + e2_1 == e1 = + Just (BinOp (SMax it) e2_2 e1, v2_cs) + | Var v2 <- e2, + Just (BasicOp (BinOp (SMax _) e2_1 e2_2), v2_cs) <- defOf v2, + e2_2 == e1 = + Just (BinOp (SMax it) e2_1 e1, v2_cs) +simplifyBinOp _ _ _ = Nothing + +constRes :: PrimValue -> Maybe (BasicOp, Certificates) +constRes = Just . (,mempty) . SubExp . Constant + +subExpRes :: SubExp -> Maybe (BasicOp, Certificates) +subExpRes = Just . (,mempty) . SubExp + +simplifyUnOp :: SimpleRule lore +simplifyUnOp _ _ (UnOp op (Constant v)) = + constRes =<< doUnOp op v +simplifyUnOp defOf _ (UnOp Not (Var v)) + | Just (BasicOp (UnOp Not v2), v_cs) <- defOf v = + Just (SubExp v2, v_cs) +simplifyUnOp _ _ _ = + Nothing + +simplifyConvOp :: SimpleRule lore +simplifyConvOp _ _ (ConvOp op (Constant v)) = + constRes =<< doConvOp op v +simplifyConvOp _ _ (ConvOp op se) + | (from, to) <- convOpType op, + from == to = + subExpRes se +simplifyConvOp lookupVar _ (ConvOp (SExt t2 t1) (Var v)) + | Just (BasicOp (ConvOp (SExt t3 _) se), v_cs) <- lookupVar v, + t2 >= t3 = + Just (ConvOp (SExt t3 t1) se, v_cs) +simplifyConvOp lookupVar _ (ConvOp (ZExt t2 t1) (Var v)) + | Just (BasicOp (ConvOp (ZExt t3 _) se), v_cs) <- lookupVar v, + t2 >= t3 = + Just (ConvOp (ZExt t3 t1) se, v_cs) +simplifyConvOp lookupVar _ (ConvOp (SIToFP t2 t1) (Var v)) + | Just (BasicOp (ConvOp (SExt t3 _) se), v_cs) <- lookupVar v, + t2 >= t3 = + Just (ConvOp (SIToFP t3 t1) se, v_cs) +simplifyConvOp lookupVar _ (ConvOp (UIToFP t2 t1) (Var v)) + | Just (BasicOp (ConvOp (ZExt t3 _) se), v_cs) <- lookupVar v, + t2 >= t3 = + Just (ConvOp (UIToFP t3 t1) se, v_cs) +simplifyConvOp lookupVar _ (ConvOp (FPConv t2 t1) (Var v)) + | Just (BasicOp (ConvOp (FPConv t3 _) se), v_cs) <- lookupVar v, + t2 >= t3 = + Just (ConvOp (FPConv t3 t1) se, v_cs) +simplifyConvOp _ _ _ = + Nothing + +-- If expression is true then just replace assertion. +simplifyAssert :: SimpleRule lore +simplifyAssert _ _ (Assert (Constant (BoolValue True)) _ _) = + constRes Checked +simplifyAssert _ _ _ = + Nothing + +simplifyIdentityReshape :: SimpleRule lore +simplifyIdentityReshape _ seType (Reshape newshape v) + | Just t <- seType $ Var v, + newDims newshape == arrayDims t -- No-op reshape. + = + subExpRes $ Var v +simplifyIdentityReshape _ _ _ = Nothing + +simplifyReshapeReshape :: SimpleRule lore +simplifyReshapeReshape defOf _ (Reshape newshape v) + | Just (BasicOp (Reshape oldshape v2), v_cs) <- defOf v = + Just (Reshape (fuseReshape oldshape newshape) v2, v_cs) +simplifyReshapeReshape _ _ _ = Nothing + +simplifyReshapeScratch :: SimpleRule lore +simplifyReshapeScratch defOf _ (Reshape newshape v) + | Just (BasicOp (Scratch bt _), v_cs) <- defOf v = + Just (Scratch bt $ newDims newshape, v_cs) +simplifyReshapeScratch _ _ _ = Nothing + +simplifyReshapeReplicate :: SimpleRule lore +simplifyReshapeReplicate defOf seType (Reshape newshape v) + | Just (BasicOp (Replicate _ se), v_cs) <- defOf v, + Just oldshape <- arrayShape <$> seType se, + shapeDims oldshape `isSuffixOf` newDims newshape = + let new = + take (length newshape - shapeRank oldshape) $ + newDims newshape + in Just (Replicate (Shape new) se, v_cs) +simplifyReshapeReplicate _ _ _ = Nothing + +simplifyReshapeIota :: SimpleRule lore +simplifyReshapeIota defOf _ (Reshape newshape v) + | Just (BasicOp (Iota _ offset stride it), v_cs) <- defOf v, + [n] <- newDims newshape = + Just (Iota n offset stride it, v_cs) +simplifyReshapeIota _ _ _ = Nothing + +improveReshape :: SimpleRule lore +improveReshape _ seType (Reshape newshape v) + | Just t <- seType $ Var v, + newshape' <- informReshape (arrayDims t) newshape, + newshape' /= newshape = + Just (Reshape newshape' v, mempty) +improveReshape _ _ _ = Nothing + +-- | If we are copying a scratch array (possibly indirectly), just turn it into a scratch by +-- itself. +copyScratchToScratch :: SimpleRule lore +copyScratchToScratch defOf seType (Copy src) = do + t <- seType $ Var src + if isActuallyScratch src + then Just (Scratch (elemType t) (arrayDims t), mempty) + else Nothing + where + isActuallyScratch v = + case asBasicOp . fst =<< defOf v of + Just Scratch {} -> True + Just (Rearrange _ v') -> isActuallyScratch v' + Just (Reshape _ v') -> isActuallyScratch v' + _ -> False +copyScratchToScratch _ _ _ = + Nothing + +simpleRules :: [SimpleRule lore] +simpleRules = + [ simplifyBinOp, + simplifyCmpOp, + simplifyUnOp, + simplifyConvOp, + simplifyAssert, + copyScratchToScratch, + simplifyIdentityReshape, + simplifyReshapeReshape, + simplifyReshapeScratch, + simplifyReshapeReplicate, + simplifyReshapeIota, + improveReshape + ] + +-- | Try to simplify the given 'BasicOp', returning a new 'BasicOp' +-- and certificates that it must depend on. +{-# NOINLINE applySimpleRules #-} +applySimpleRules :: + VarLookup lore -> + TypeLookup -> + BasicOp -> + Maybe (BasicOp, Certificates) +applySimpleRules defOf seType op = + msum [rule defOf seType op | rule <- simpleRules] 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 @@ -8,12 +8,13 @@ import Control.Monad.Reader import Control.Monad.State -import Data.List (foldl') import qualified Data.Map.Strict as M import Data.Maybe (mapMaybe) import qualified Data.Sequence as Seq import Futhark.IR.Kernels import Futhark.MonadFreshNames +import Futhark.Optimise.BlkRegTiling +import Futhark.Optimise.TileLoops.Shared import Futhark.Pass import Futhark.Tools import Futhark.Transform.Rename @@ -30,8 +31,6 @@ runState $ runReaderT (optimiseStms stms) scope -type TileM = ReaderT (Scope Kernels) (State VNameSource) - optimiseBody :: Body Kernels -> TileM (Body Kernels) optimiseBody (Body () stms res) = Body () <$> optimiseStms stms <*> pure res @@ -42,11 +41,17 @@ mconcat <$> mapM optimiseStm (stmsToList stms) optimiseStm :: Stm Kernels -> TileM (Stms Kernels) -optimiseStm (Let pat aux (Op (SegOp (SegMap lvl@SegThread {} space ts kbody)))) = do - (host_stms, (lvl', space', kbody')) <- tileInKernelBody mempty initial_variance lvl space ts kbody - return $ - host_stms - <> oneStm (Let pat aux $ Op $ SegOp $ SegMap lvl' space' ts kbody') +optimiseStm stm@(Let pat aux (Op (SegOp (SegMap lvl@SegThread {} space ts kbody)))) = do + res3dtiling <- doRegTiling3D stm + case res3dtiling of + Just (extra_bnds, stmt') -> return (extra_bnds <> oneStm stmt') + Nothing -> do + blkRegTiling_res <- mmBlkRegTiling stm + case blkRegTiling_res of + Just (extra_bnds, stmt') -> return (extra_bnds <> oneStm stmt') + Nothing -> do + (host_stms, (lvl', space', kbody')) <- tileInKernelBody mempty initial_variance lvl space ts kbody + return $ host_stms <> oneStm (Let pat aux $ Op $ SegOp $ SegMap lvl' space' ts kbody') where initial_variance = M.map mempty $ scopeOfSegSpace space optimiseStm (Let pat aux e) = @@ -65,12 +70,12 @@ tileInKernelBody branch_variant initial_variance lvl initial_kspace ts kbody | Just kbody_res <- mapM isSimpleResult $ kernelBodyResult kbody = do maybe_tiled <- - tileInBody branch_variant mempty initial_variance lvl initial_kspace ts $ + tileInBody branch_variant initial_variance lvl initial_kspace ts $ Body () (kernelBodyStms kbody) kbody_res case maybe_tiled of Just (host_stms, tiling, tiledBody) -> do (res', stms') <- - runBinder $ mapM (tilingTileReturns tiling) =<< tiledBody mempty + runBinder $ mapM (tilingTileReturns tiling) =<< tiledBody mempty mempty return ( host_stms, ( tilingLevel tiling, @@ -88,14 +93,13 @@ tileInBody :: Names -> - Names -> VarianceTable -> SegLevel -> SegSpace -> [Type] -> Body Kernels -> TileM (Maybe (Stms Kernels, Tiling, TiledBody)) -tileInBody branch_variant private initial_variance initial_lvl initial_space res_ts (Body () initial_kstms stms_res) = +tileInBody branch_variant initial_variance initial_lvl initial_space res_ts (Body () initial_kstms stms_res) = descend mempty $ stmsToList initial_kstms where variance = varianceInStms initial_variance initial_kstms @@ -114,7 +118,7 @@ (prestms', poststms') <- preludeToPostlude variance prestms stm_to_tile (stmsFromList poststms), used <- freeIn stm_to_tile <> freeIn poststms' <> freeIn stms_res = - Just . injectPrelude initial_space private variance prestms' used + Just . injectPrelude initial_space variance prestms' used <$> tileGeneric (tiling2d $ reverse $ zip top_gtids_rev top_kdims_rev) initial_lvl @@ -136,7 +140,7 @@ (prestms', poststms') <- preludeToPostlude variance prestms stm_to_tile (stmsFromList poststms), used <- freeIn stm_to_tile <> freeIn poststms' <> freeIn stms_res = - Just . injectPrelude initial_space private variance prestms' used + Just . injectPrelude initial_space variance prestms' used <$> tileGeneric (tiling1d $ reverse top_space_rev) initial_lvl @@ -161,13 +165,11 @@ (namesToList (freeIn bound)) ) merge_params = map fst merge - private' = namesFromList $ map paramName merge_params maybe_tiled <- localScope (M.insert i (IndexName it) $ scopeOfFParams merge_params) $ tileInBody branch_variant' - private' variance initial_lvl initial_space @@ -284,25 +286,20 @@ -- considered thread-local. injectPrelude :: SegSpace -> - Names -> VarianceTable -> Stms Kernels -> Names -> (Stms Kernels, Tiling, TiledBody) -> (Stms Kernels, Tiling, TiledBody) -injectPrelude initial_space private variance prestms used (host_stms, tiling, tiledBody) = +injectPrelude initial_space variance prestms used (host_stms, tiling, tiledBody) = (host_stms, tiling, tiledBody') where - private' = - private - <> namesFromList - ( map fst $ - filter (`notElem` unSegSpace (tilingSpace tiling)) $ - unSegSpace initial_space - ) - - tiledBody' privstms = do - let ( invariant_prestms, + tiledBody' private privstms = do + let nontiled = (`notElem` unSegSpace (tilingSpace tiling)) + private' = + private + <> namesFromList (map fst (filter nontiled $ unSegSpace initial_space)) + ( invariant_prestms, precomputed_variant_prestms, recomputed_variant_prestms ) = @@ -322,7 +319,7 @@ PrivStms recomputed_variant_prestms $ mkReadPreludeValues prelude_arrs live_set - tiledBody (prelude_privstms <> privstms) + tiledBody private' (prelude_privstms <> privstms) tileDoLoop :: SegSpace -> @@ -341,10 +338,7 @@ Result -> TileM (Stms Kernels, Tiling, TiledBody) tileDoLoop initial_space variance prestms used_in_body (host_stms, tiling, tiledBody) res_ts pat aux merge i it bound poststms poststms_res = do - let prestms_used = - used_in_body - <> freeIn poststms - <> freeIn poststms_res + let prestms_used = used_in_body <> freeIn poststms <> freeIn poststms_res ( invariant_prestms, precomputed_variant_prestms, recomputed_variant_prestms @@ -358,7 +352,7 @@ merge_scope = M.insert i (IndexName it) $ scopeOfFParams mergeparams - tiledBody' privstms = localScope (scopeOf host_stms <> merge_scope) $ do + tiledBody' private privstms = localScope (scopeOf host_stms <> merge_scope) $ do addStms invariant_prestms let live_set = @@ -400,10 +394,16 @@ Index (paramName from) $ fullSlice (paramType from) slice + private' = + private <> namesFromList (map paramName mergeparams ++ map paramName mergeparams') + + privstms' = + PrivStms mempty indexMergeParams <> privstms <> inloop_privstms + loopbody' <- runBodyBinder $ resultBody . map Var - <$> tiledBody (PrivStms mempty indexMergeParams <> privstms <> inloop_privstms) + <$> tiledBody private' privstms' accs' <- letTupExp "tiled_inside_loop" $ DoLoop [] merge' (ForLoop i it bound []) loopbody' @@ -626,7 +626,7 @@ addStms poststms return poststms_res -type TiledBody = PrivStms -> Binder Kernels [VName] +type TiledBody = Names -> PrivStms -> Binder Kernels [VName] tileGeneric :: DoTiling gtids kdims -> @@ -648,8 +648,8 @@ where (red_comm, red_lam, red_nes, map_lam) = form - tiledBody :: Tiling -> PrivStms -> Binder Kernels [VName] - tiledBody tiling privstms = do + tiledBody :: Tiling -> Names -> PrivStms -> Binder Kernels [VName] + tiledBody tiling _private privstms = do let tile_shape = tilingTileShape tiling num_whole_tiles <- tilingNumWholeTiles tiling @@ -981,30 +981,6 @@ then Just $ InputTile [1, 0] arr else Just $ InputDontTile arr -segMap2D :: - String -> - SegLevel -> - ResultManifest -> - (SubExp, SubExp) -> - ((VName, VName) -> Binder Kernels [SubExp]) -> - Binder Kernels [VName] -segMap2D desc lvl manifest (dim_x, dim_y) f = do - ltid_x <- newVName "ltid_x" - ltid_y <- newVName "ltid_y" - ltid_flat <- newVName "ltid_flat" - let space = SegSpace ltid_flat [(ltid_x, dim_x), (ltid_y, dim_y)] - - ((ts, res), stms) <- runBinder $ do - res <- f (ltid_x, ltid_y) - ts <- mapM subExpType res - return (ts, res) - Body _ stms' res' <- renameBody $ mkBody stms res - - letTupExp desc $ - Op $ - SegOp $ - SegMap lvl space ts $ KernelBody () stms' $ map (Returns manifest) res' - -- Reconstruct the original gtids from group and local IDs. reconstructGtids2D :: SubExp -> @@ -1279,21 +1255,3 @@ tilingLevel = lvl, tilingSpace = space } - --- | The variance table keeps a mapping from a variable name --- (something produced by a 'Stm') to the kernel thread indices --- that name depends on. If a variable is not present in this table, --- that means it is bound outside the kernel (and so can be considered --- invariant to all dimensions). -type VarianceTable = M.Map VName Names - -varianceInStms :: VarianceTable -> Stms Kernels -> VarianceTable -varianceInStms = foldl varianceInStm - -varianceInStm :: VarianceTable -> Stm Kernels -> VarianceTable -varianceInStm variance bnd = - foldl' add variance $ patternNames $ stmPattern bnd - 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 variance) $ namesToList (freeIn bnd) diff --git a/src/Futhark/Optimise/TileLoops/Shared.hs b/src/Futhark/Optimise/TileLoops/Shared.hs new file mode 100644 --- /dev/null +++ b/src/Futhark/Optimise/TileLoops/Shared.hs @@ -0,0 +1,160 @@ +module Futhark.Optimise.TileLoops.Shared + ( TileM, + segMap2D, + segMap3D, + segScatter2D, + VarianceTable, + varianceInStms, + isTileableRedomap, + ) +where + +import Control.Monad.Reader +import Control.Monad.State +import Data.List (foldl', zip4) +import qualified Data.Map as M +import Futhark.IR.Kernels +import Futhark.MonadFreshNames +import Futhark.Tools +import Futhark.Transform.Rename + +type TileM = ReaderT (Scope Kernels) (State VNameSource) + +segMap2D :: + String -> -- desc + SegLevel -> -- lvl + ResultManifest -> -- manifest + (SubExp, SubExp) -> -- (dim_x, dim_y) + ( (VName, VName) -> -- f + Binder Kernels [SubExp] + ) -> + Binder Kernels [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" + let segspace = SegSpace ltid_flat [(ltid_yy, dim_y), (ltid_xx, dim_x)] + + ((ts, res), stms) <- runBinder $ do + res <- f (ltid_yy, ltid_xx) + ts <- mapM subExpType res + return (ts, res) + + letTupExp desc <=< renameExp $ + Op $ + SegOp $ + SegMap lvl segspace ts $ KernelBody () stms $ map (Returns manifest) res + +segMap3D :: + String -> -- desc + SegLevel -> -- lvl + ResultManifest -> -- manifest + (SubExp, SubExp, SubExp) -> -- (dim_z, dim_y, dim_x) + ( (VName, VName, VName) -> -- f + Binder Kernels [SubExp] + ) -> + Binder Kernels [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" + let segspace = SegSpace ltid_flat [(ltid_z, dim_z), (ltid_y, dim_y), (ltid_x, dim_x)] + + ((ts, res), stms) <- runBinder $ do + res <- f (ltid_z, ltid_y, ltid_x) + ts <- mapM subExpType res + return (ts, res) + + letTupExp desc <=< renameExp $ + Op $ + SegOp $ + SegMap lvl segspace ts $ KernelBody () stms $ map (Returns manifest) res + +segScatter2D :: + String -> -- desc + SubExp -> -- arr_size + VName -> + SegLevel -> -- lvl + (SubExp, SubExp) -> -- (dim_y, dim_x) + ((VName, VName) -> Binder Kernels (SubExp, SubExp)) -> -- f + Binder Kernels [VName] +segScatter2D desc arr_size updt_arr lvl (dim_x, dim_y) f = do + ltid_x <- newVName "ltid_x" + ltid_y <- newVName "ltid_y" + ltid_flat <- newVName "ltid_flat" + let segspace = SegSpace ltid_flat [(ltid_x, dim_x), (ltid_y, dim_y)] + + ((t_v, res_v, res_i), stms) <- runBinder $ do + (res_v, res_i) <- f (ltid_x, ltid_y) + t_v <- subExpType res_v + return (t_v, res_v, res_i) + + let ret = WriteReturns [arr_size] updt_arr [([DimFix res_i], res_v)] + let body = KernelBody () stms [ret] + + letTupExp desc <=< renameExp $ Op $ SegOp $ SegMap lvl segspace [t_v] body + +-- | The variance table keeps a mapping from a variable name +-- (something produced by a 'Stm') to the kernel thread indices +-- that name depends on. If a variable is not present in this table, +-- that means it is bound outside the kernel (and so can be considered +-- invariant to all dimensions). +type VarianceTable = M.Map VName Names + +isTileableRedomap :: + Stm Kernels -> + Maybe + ( SubExp, + [VName], + (Commutativity, Lambda Kernels, [SubExp], Lambda Kernels) + ) +isTileableRedomap stm + | Op (OtherOp (Screma w form arrs)) <- stmExp stm, + Just (reds, map_lam) <- isRedomapSOAC form, + Reduce red_comm red_lam red_nes <- singleReduce reds, + all (primType . rowType . paramType) $ lambdaParams red_lam, + all (primType . rowType . paramType) $ lambdaParams map_lam, + lambdaReturnType map_lam == lambdaReturnType red_lam, -- No mapout arrays. + not (null arrs), + all primType $ lambdaReturnType map_lam, + all (primType . paramType) $ lambdaParams map_lam = + Just (w, arrs, (red_comm, red_lam, red_nes, map_lam)) + | otherwise = + Nothing + +defVarianceInStm :: VarianceTable -> Stm Kernels -> VarianceTable +defVarianceInStm variance bnd = + foldl' add variance $ patternNames $ stmPattern bnd + 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 variance) $ namesToList (freeIn bnd) + +-- just in case you need the Screma being treated differently than +-- by default; previously Cosmin had to enhance it when dealing with stream. +varianceInStm :: VarianceTable -> Stm Kernels -> VarianceTable +varianceInStm v0 bnd@(Let _ _ (Op (OtherOp Screma {}))) + | Just (_, arrs, (_, red_lam, red_nes, map_lam)) <- isTileableRedomap bnd = + let v = defVarianceInStm v0 bnd + red_ps = lambdaParams red_lam + map_ps = lambdaParams map_lam + card_red = length red_nes + acc_lam_f = take (card_red `quot` 2) red_ps + arr_lam_f = drop (card_red `quot` 2) red_ps + stm_lam = bodyStms (lambdaBody map_lam) <> bodyStms (lambdaBody red_lam) + + f vacc (v_a, v_fm, v_fr_acc, v_fr_var) = + let vrc = oneName v_a <> M.findWithDefault mempty v_a vacc + vacc' = M.insert v_fm vrc vacc + vrc' = oneName v_fm <> vrc + in M.insert v_fr_acc (oneName v_fr_var <> vrc') $ M.insert v_fr_var vrc' vacc' + + v' = + foldl' f v $ + zip4 arrs (map paramName map_ps) (map paramName acc_lam_f) (map paramName arr_lam_f) + in varianceInStms v' stm_lam +varianceInStm v0 bnd = defVarianceInStm v0 bnd + +varianceInStms :: VarianceTable -> Stms Kernels -> VarianceTable +varianceInStms = foldl' varianceInStm 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 @@ -307,10 +307,12 @@ stmsToList $ get_stms body in (set_stms (stmsFromList stms) body, allocs) -expandable :: Space -> Bool +expandable, notScalar :: Space -> Bool expandable (Space "local") = False expandable ScalarSpace {} = False expandable _ = True +notScalar ScalarSpace {} = False +notScalar _ = True extractStmAllocations :: SegLevel -> @@ -319,7 +321,12 @@ Stm KernelsMem -> Writer Extraction (Maybe (Stm KernelsMem)) extractStmAllocations lvl bound_outside bound_kernel (Let (Pattern [] [patElem]) _ (Op (Alloc size space))) - | expandable space && expandableSize size || boundInKernel size = do + | expandable space && expandableSize size + -- FIXME: the '&& notScalar space' part is a hack because we + -- don't otherwise hoist the sizes out far enough, and we + -- promise to be super-duper-careful about not having variant + -- scalar allocations. + || (boundInKernel size && notScalar space) = do tell $ M.singleton (patElemName patElem) (lvl, size, space) return Nothing where diff --git a/src/Futhark/Pass/ExplicitAllocations.hs b/src/Futhark/Pass/ExplicitAllocations.hs --- a/src/Futhark/Pass/ExplicitAllocations.hs +++ b/src/Futhark/Pass/ExplicitAllocations.hs @@ -574,6 +574,9 @@ Space -> VName -> AllocM fromlore tolore (SubExp, ExtIxFun, [TPrimExp Int64 VName], VName) +existentializeArray ScalarSpace {} v = do + (mem', ixfun) <- lookupArraySummary v + return (Var v, fmap (fmap Free) ixfun, mempty, mem') existentializeArray space v = do (mem', ixfun) <- lookupArraySummary v sp <- lookupMemSpace mem' diff --git a/src/Futhark/Pass/ExtractKernels.hs b/src/Futhark/Pass/ExtractKernels.hs --- a/src/Futhark/Pass/ExtractKernels.hs +++ b/src/Futhark/Pass/ExtractKernels.hs @@ -163,6 +163,7 @@ import Control.Monad.Identity import Control.Monad.RWS.Strict import Control.Monad.Reader +import Data.Function ((&)) import Data.Maybe import qualified Futhark.IR.Kernels as Out import Futhark.IR.Kernels.Kernel @@ -278,7 +279,7 @@ bodyStms body -- XXX - our notion of balancing is probably still too naive. - unbalancedStm bound (Op (Stream w _ _ _)) = + unbalancedStm bound (Op (Stream w _ _ _ _)) = w `subExpBound` bound unbalancedStm bound (Op (Screma w _ _)) = w `subExpBound` bound @@ -375,7 +376,7 @@ Scan scan_lam nes <- singleScan scans, Just do_iswim <- iswim res_pat w scan_lam $ zip nes arrs = do types <- asksScope scopeForSOACs - transformStms path =<< (stmsToList . snd <$> runBinderT (certifying cs do_iswim) types) + transformStms path . stmsToList . snd =<< runBinderT (certifying cs do_iswim) types | Just (scans, map_lam) <- isScanomapSOAC form = runBinder_ $ do scan_ops <- forM scans $ \(Scan scan_lam nes) -> do (scan_lam', nes', shape) <- determineReduceOp scan_lam nes @@ -383,10 +384,8 @@ return $ SegBinOp Noncommutative scan_lam'' nes' shape let map_lam_sequential = soacsLambdaToKernels map_lam lvl <- segThreadCapped [w] "segscan" $ NoRecommendation SegNoVirt - addStms - =<< ( fmap (certify cs) - <$> segScan lvl res_pat w scan_ops map_lam_sequential arrs [] [] - ) + addStms . fmap (certify cs) + =<< segScan lvl res_pat w scan_ops map_lam_sequential arrs [] [] -- We are only willing to generate code for scanomaps that do not -- involve array accumulators, and do not have parallelism in their @@ -428,18 +427,21 @@ return $ SegBinOp comm' red_lam'' nes' shape let map_lam_sequential = soacsLambdaToKernels map_lam lvl <- segThreadCapped [w] "segred" $ NoRecommendation SegNoVirt - addStms - =<< ( fmap (certify cs) - <$> nonSegRed lvl pat w red_ops map_lam_sequential arrs - ) + addStms . fmap (certify cs) + =<< nonSegRed lvl pat w red_ops map_lam_sequential arrs outerParallelBody = renameBody =<< (mkBody <$> paralleliseOuter <*> pure (map Var (patternNames pat))) paralleliseInner path' = do - (mapbnd, redbnd) <- redomapToMapAndReduce pat (w, comm', red_lam, map_lam, nes, arrs) - transformStms path' [certify cs mapbnd, certify cs redbnd] + (mapstm, redstm) <- + redomapToMapAndReduce pat (w, comm', red_lam, map_lam, nes, arrs) + types <- asksScope scopeForSOACs + transformStms path' . stmsToList <=< (`runBinderT_` types) $ do + (_, stms) <- + simplifyStms (stmsFromList [certify cs mapstm, certify cs redstm]) + addStms stms where comm' | commutativeLambda red_lam = Commutative @@ -464,14 +466,14 @@ -- Streams can be handled in two different ways - either we -- sequentialise the body or we keep it parallel and distribute. -transformStm path (Let pat aux@(StmAux cs _ _) (Op (Stream w (Parallel _ _ _ []) map_fun arrs))) +transformStm path (Let pat aux@(StmAux cs _ _) (Op (Stream w Parallel {} map_fun [] arrs))) | not ("sequential_inner" `inAttrs` stmAuxAttrs aux) = do -- No reduction part. Remove the stream and leave the body -- parallel. It will be distributed. types <- asksScope scopeForSOACs - transformStms path - =<< (stmsToList . snd <$> runBinderT (certifying cs $ sequentialStreamWholeArray pat w [] map_fun arrs) types) -transformStm path (Let pat aux@(StmAux cs _ _) (Op (Stream w (Parallel o comm red_fun nes) fold_fun arrs))) + transformStms path . stmsToList . snd + =<< runBinderT (certifying cs $ sequentialStreamWholeArray pat w [] map_fun arrs) types +transformStm path (Let pat aux@(StmAux cs _ _) (Op (Stream w (Parallel o comm red_fun) fold_fun nes arrs))) | "sequential_inner" `inAttrs` stmAuxAttrs aux = paralleliseOuter path | otherwise = do @@ -534,8 +536,8 @@ paralleliseInner path' = do types <- asksScope scopeForSOACs - transformStms path' . fmap (certify cs) - =<< (stmsToList . snd <$> runBinderT (sequentialStreamWholeArray pat w nes fold_fun arrs) types) + transformStms path' . fmap (certify cs) . stmsToList . snd + =<< runBinderT (sequentialStreamWholeArray pat w nes fold_fun arrs) types innerParallelBody path' = renameBody @@ -550,29 +552,36 @@ scope <- asksScope scopeForSOACs transformStms path . map (certify cs) . stmsToList . snd =<< runBinderT (dissectScrema pat w form arrs) scope -transformStm path (Let pat _ (Op (Stream w (Sequential nes) fold_fun arrs))) = do +transformStm path (Let pat _ (Op (Stream w Sequential fold_fun nes arrs))) = do -- Remove the stream and leave the body parallel. It will be -- distributed. types <- asksScope scopeForSOACs - transformStms path - =<< ( stmsToList . snd - <$> runBinderT (sequentialStreamWholeArray pat w nes fold_fun arrs) types - ) + transformStms path . stmsToList . snd + =<< runBinderT (sequentialStreamWholeArray pat w nes fold_fun arrs) types transformStm _ (Let pat (StmAux cs _ _) (Op (Scatter w lam ivs as))) = runBinder_ $ do let lam' = soacsLambdaToKernels lam write_i <- newVName "write_i" let (as_ws, as_ns, as_vs) = unzip3 as - (i_res, v_res) = splitAt (sum as_ns) $ bodyResult $ lambdaBody lam' + indexes = zipWith (*) as_ns $ map length as_ws + (i_res, v_res) = splitAt (sum indexes) $ bodyResult $ lambdaBody lam' kstms = bodyStms $ lambdaBody lam' krets = do - (a_w, a, is_vs) <- zip3 as_ws as_vs $ chunks as_ns $ zip i_res v_res - return $ WriteReturns [a_w] a [([DimFix i], v) | (i, v) <- is_vs] + (a_w, a, is_vs) <- + zip (chunks (concat $ zipWith (\ws n -> replicate n $ length ws) as_ws as_ns) i_res) v_res + & chunks as_ns + & zip3 as_ws as_vs + return $ WriteReturns (shapeDims a_w) a [(map DimFix is, v) | (is, v) <- is_vs] body = KernelBody () kstms krets inputs = do (p, p_a) <- zip (lambdaParams lam') ivs return $ KernelInput (paramName p) (paramType p) p_a [Var write_i] (kernel, stms) <- - mapKernel segThreadCapped [(write_i, w)] inputs (map rowType $ patternTypes pat) body + mapKernel + segThreadCapped + [(write_i, w)] + inputs + (zipWith (stripArray . length) as_ws $ patternTypes pat) + body certifying cs $ do addStms stms letBind pat $ Op $ SegOp kernel @@ -621,7 +630,7 @@ max (bodyInterest tbody) (bodyInterest fbody) | Op (Screma w (ScremaForm _ _ lam') _) <- stmExp stm = zeroIfTooSmall w + bodyInterest (lambdaBody lam') - | Op (Stream _ (Sequential _) lam' _) <- stmExp stm = + | Op (Stream _ Sequential lam' _ _) <- stmExp stm = bodyInterest $ lambdaBody lam' | otherwise = 0 diff --git a/src/Futhark/Pass/ExtractKernels/DistributeNests.hs b/src/Futhark/Pass/ExtractKernels/DistributeNests.hs --- a/src/Futhark/Pass/ExtractKernels/DistributeNests.hs +++ b/src/Futhark/Pass/ExtractKernels/DistributeNests.hs @@ -7,6 +7,7 @@ {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} +{-# OPTIONS_GHC -Wno-overlapping-patterns -Wno-incomplete-patterns -Wno-incomplete-uni-patterns -Wno-incomplete-record-updates #-} module Futhark.Pass.ExtractKernels.DistributeNests ( MapLoop (..), @@ -39,6 +40,7 @@ import Control.Monad.Reader import Control.Monad.Trans.Maybe import Control.Monad.Writer.Strict +import Data.Function ((&)) import Data.List (find, partition, tails) import qualified Data.Map as M import Data.Maybe @@ -337,7 +339,7 @@ distributeMapBodyStms orig_acc = distribute <=< onStms orig_acc . stmsToList where onStms acc [] = return acc - onStms acc (Let pat (StmAux cs _ _) (Op (Stream w (Sequential accs) lam arrs)) : stms) = do + onStms acc (Let pat (StmAux cs _ _) (Op (Stream w Sequential lam accs arrs)) : stms) = do types <- asksScope scopeForSOACs stream_stms <- snd <$> runBinderT (sequentialStreamWholeArray pat w accs lam arrs) types @@ -626,7 +628,7 @@ lambdaReturnType = [Prim int64, et], lambdaBody = mkBody mempty [i, v] } - maybeDistributeStm (Let pat aux $ Op $ Scatter (intConst Int64 1) lam [] [(w, 1, arr)]) acc + maybeDistributeStm (Let pat aux $ Op $ Scatter (intConst Int64 1) lam [] [(Shape [w], 1, arr)]) acc where amortises DoLoop {} = True amortises Op {} = True @@ -758,7 +760,7 @@ SubExp -> Lambda lore -> [VName] -> - [(SubExp, Int, VName)] -> + [(Shape, Int, VName)] -> DistNestT lore m (Stms lore) segmentedScatterKernel nest perm scatter_pat cs scatter_w lam ivs dests = do -- We replicate some of the checking done by 'isSegmentedOp', but @@ -775,6 +777,7 @@ (ispace, kernel_inps) <- flatKernel nest' let (as_ws, as_ns, as) = unzip3 dests + indexes = zipWith (*) as_ns $ map length as_ws -- The input/output arrays ('as') _must_ correspond to some kernel -- input, or else the original nested scatter would have been @@ -786,8 +789,8 @@ let rts = concatMap (take 1) $ chunks as_ns $ - drop (sum as_ns) $ lambdaReturnType lam - (is, vs) = splitAt (sum as_ns) $ bodyResult $ lambdaBody lam + drop (sum indexes) $ lambdaReturnType lam + (is, vs) = splitAt (sum indexes) $ bodyResult $ lambdaBody lam -- Maybe add certificates to the indices. (is', k_body_stms) <- runBinder $ do @@ -798,9 +801,11 @@ else certifying cs $ letSubExp "scatter_i" $ BasicOp $ SubExp i let k_body = - KernelBody () k_body_stms $ - map (inPlaceReturn ispace) $ - zip3 as_ws as_inps $ chunks as_ns $ zip is' vs + zip (chunks (concat $ zipWith (\ws n -> replicate n $ length ws) as_ws as_ns) is') vs + & chunks as_ns + & zip3 (map shapeDims as_ws) as_inps + & map (inPlaceReturn ispace) + & KernelBody () k_body_stms (k, k_bnds) <- mapKernel mk_lvl ispace kernel_inps rts k_body @@ -820,9 +825,9 @@ inPlaceReturn ispace (aw, inp, is_vs) = WriteReturns - (init ws ++ [aw]) + (init ws ++ aw) (kernelInputArray inp) - [(map DimFix $ map Var (init gtids) ++ [i], v) | (i, v) <- is_vs] + [(map DimFix $ map Var (init gtids) ++ is, v) | (is, v) <- is_vs] where (gtids, ws) = unzip ispace @@ -985,7 +990,7 @@ =<< segHist lvl orig_pat hist_w ispace inputs' ops' lam arrs determineReduceOp :: - (MonadBinder m, Lore m ~ lore) => + MonadBinder m => Lambda SOACS -> [SubExp] -> m (Lambda SOACS, [SubExp], Shape) diff --git a/src/Futhark/Pass/ExtractKernels/Distribution.hs b/src/Futhark/Pass/ExtractKernels/Distribution.hs --- a/src/Futhark/Pass/ExtractKernels/Distribution.hs +++ b/src/Futhark/Pass/ExtractKernels/Distribution.hs @@ -3,6 +3,7 @@ {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE LambdaCase #-} {-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TupleSections #-} {-# LANGUAGE TypeFamilies #-} @@ -356,7 +357,8 @@ distributionBodyFromStms targets $ oneStm bnd createKernelNest :: - (MonadFreshNames m, HasScope t m) => + forall lore m. + (MonadFreshNames m, HasScope lore m) => Nestings -> DistributionBody -> m (Maybe (Targets, KernelNest)) @@ -379,7 +381,6 @@ (== mempty) . namesIntersection bound_in_nest . freeIn . arrayDims distributeAtNesting :: - (HasScope t m, MonadFreshNames m) => Nesting -> PatternT Type -> (LoopNesting -> KernelNest, Names) -> @@ -462,7 +463,6 @@ ) recurse :: - (HasScope t m, MonadFreshNames m) => [(Nesting, Target)] -> MaybeT m (KernelNest, Names, Targets) recurse [] = diff --git a/src/Futhark/Pass/ExtractKernels/Intragroup.hs b/src/Futhark/Pass/ExtractKernels/Intragroup.hs --- a/src/Futhark/Pass/ExtractKernels/Intragroup.hs +++ b/src/Futhark/Pass/ExtractKernels/Intragroup.hs @@ -10,6 +10,7 @@ import Control.Monad.Identity import Control.Monad.RWS import Control.Monad.Trans.Maybe +import Data.Function ((&)) import qualified Data.Map.Strict as M import qualified Data.Set as S import Futhark.Analysis.PrimExp.Convert @@ -260,7 +261,7 @@ certifying (stmAuxCerts aux) $ addStms =<< segHist lvl' pat w [] [] ops' bucket_fun' arrs parallelMin [w] - Op (Stream w (Sequential accs) lam arrs) + Op (Stream w Sequential lam accs arrs) | chunk_size_param : _ <- lambdaParams lam -> do types <- asksScope castScope ((), stream_bnds) <- @@ -276,10 +277,14 @@ let lam' = soacsLambdaToKernels lam (dests_ws, dests_ns, dests_vs) = unzip3 dests - (i_res, v_res) = splitAt (sum dests_ns) $ bodyResult $ lambdaBody lam' + indexes = zipWith (*) dests_ns $ map length dests_ws + (i_res, v_res) = splitAt (sum indexes) $ bodyResult $ lambdaBody lam' krets = do - (a_w, a, is_vs) <- zip3 dests_ws dests_vs $ chunks dests_ns $ zip i_res v_res - return $ WriteReturns [a_w] a [([DimFix i], v) | (i, v) <- is_vs] + (a_w, a, is_vs) <- + zip (chunks (concat $ zipWith (\ws n -> replicate n $ length ws) dests_ws dests_ns) i_res) v_res + & chunks dests_ns + & zip3 dests_ws dests_vs + return $ WriteReturns (shapeDims a_w) a [(map DimFix is, v) | (is, v) <- is_vs] inputs = do (p, p_a) <- zip (lambdaParams lam') ivs return $ KernelInput (paramName p) (paramType p) p_a [Var write_i] @@ -290,7 +295,7 @@ addStms $ bodyStms $ lambdaBody lam' certifying (stmAuxCerts aux) $ do - let ts = map rowType $ patternTypes pat + let ts = zipWith (stripArray . length) dests_ws $ patternTypes pat body = KernelBody () kstms krets letBind pat $ Op $ SegOp $ SegMap lvl' space ts body 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 @@ -7,6 +7,7 @@ import Control.Monad.Identity import Control.Monad.Reader import Control.Monad.State +import Data.Function ((&)) import Futhark.Analysis.Rephrase import Futhark.IR import Futhark.IR.MC @@ -318,13 +319,15 @@ Body () kstms res <- mapLambdaToBody transformBody gtid lam ivs let (dests_ws, dests_ns, dests_vs) = unzip3 dests - (i_res, v_res) = splitAt (sum dests_ns) res + indexes = zipWith (*) dests_ns $ map length dests_ws + (i_res, v_res) = splitAt (sum indexes) res rets = takeLast (length dests) $ lambdaReturnType lam kres = do (a_w, a, is_vs) <- - zip3 dests_ws dests_vs $ - chunks dests_ns $ zip i_res v_res - return $ WriteReturns [a_w] a [([DimFix i], v) | (i, v) <- is_vs] + zip (chunks (concat $ zipWith (\ws n -> replicate n $ length ws) dests_ws dests_ns) i_res) v_res + & chunks dests_ns + & zip3 dests_ws dests_vs + return $ WriteReturns (shapeDims a_w) a [(map DimFix is, v) | (is, v) <- is_vs] kbody = KernelBody () kstms kres return $ oneStm $ @@ -347,7 +350,7 @@ return $ mconcat seq_hist_stms <> oneStm (Let pat (defAux ()) $ Op $ ParOp Nothing seq_op) -transformSOAC pat attrs (Stream w (Parallel _ comm red_lam red_nes) fold_lam arrs) +transformSOAC pat attrs (Stream w (Parallel _ comm red_lam) fold_lam red_nes arrs) | not $ null red_nes = do map_lam <- unstreamLambda attrs red_nes fold_lam (seq_red_stms, seq_op) <- @@ -380,12 +383,12 @@ return $ seq_red_stms <> oneStm (Let pat (defAux ()) $ Op $ ParOp Nothing seq_op) -transformSOAC pat _ (Stream w form lam arrs) = do +transformSOAC pat _ (Stream w _ lam nes arrs) = do -- Just remove the stream and transform the resulting stms. soacs_scope <- castScope <$> askScope stream_stms <- flip runBinderT_ soacs_scope $ - sequentialStreamWholeArray pat w (getStreamAccums form) lam arrs + sequentialStreamWholeArray pat w nes lam arrs transformStms stream_stms transformProg :: Prog SOACS -> PassM (Prog MC) diff --git a/src/Futhark/Script.hs b/src/Futhark/Script.hs --- a/src/Futhark/Script.hs +++ b/src/Futhark/Script.hs @@ -1,5 +1,6 @@ {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} -- | FutharkScript is a (tiny) subset of Futhark used to write small -- expressions that are evaluated by server executables. The @futhark @@ -10,6 +11,7 @@ withScriptServer, -- * Expressions, values, and types + Func (..), Exp (..), parseExp, varsInExp, @@ -19,6 +21,7 @@ ExpValue, -- * Evaluation + EvalBuiltin, evalExp, getExpValue, evalExpToGround, @@ -28,6 +31,7 @@ where import Control.Monad.Except +import qualified Data.Binary as Bin import qualified Data.ByteString.Lazy.Char8 as LBS import Data.Char import Data.Foldable (toList) @@ -36,7 +40,6 @@ import qualified Data.Map as M import qualified Data.Set as S import qualified Data.Text as T -import qualified Data.Text.IO as T import Data.Traversable import Data.Void import Futhark.Server @@ -49,6 +52,7 @@ import System.IO import System.IO.Temp import Text.Megaparsec +import Text.Megaparsec.Char.Lexer (charLiteral) -- | Like a 'Server', but keeps a bit more state to make FutharkScript -- more convenient. @@ -61,40 +65,46 @@ counter <- newIORef 0 f $ ScriptServer server counter +-- | A function called in a 'Call' expression can be either a Futhark +-- function or a builtin function. +data Func = FuncFut EntryName | FuncBuiltin T.Text + deriving (Show) + -- | A FutharkScript expression. This is a simple AST that might not -- correspond exactly to what the user wrote (e.g. no parentheses or -- source locations). This is fine for small expressions, which is -- all this is meant for. data Exp - = Call EntryName [Exp] + = Call Func [Exp] | Const PrimValue | Tuple [Exp] | Record [(T.Text, Exp)] + | StringLit T.Text | -- | Server-side variable, *not* Futhark variable (these are -- handled in 'Call'). ServerVar TypeName VarName deriving (Show) +instance Pretty Func where + ppr (FuncFut f) = ppr f + ppr (FuncBuiltin f) = "$" <> ppr f + instance Pretty Exp where ppr = pprPrec 0 - pprPrec _ (ServerVar _ v) = "$" <> strictText v + pprPrec _ (ServerVar _ v) = "$" <> ppr v pprPrec _ (Const v) = ppr v + pprPrec _ (Call v []) = ppr v pprPrec i (Call v args) = - parensIf (i > 0) $ strictText v <+> spread (map (pprPrec 1) args) + parensIf (i > 0) $ ppr v <+> spread (map (pprPrec 1) args) pprPrec _ (Tuple vs) = parens $ commasep $ map ppr vs + pprPrec _ (StringLit s) = ppr $ show s pprPrec _ (Record m) = braces $ commasep $ map field m where field (k, v) = ppr k <> equals <> ppr v type Parser = Parsec Void T.Text -parseEntryName :: Parser EntryName -parseEntryName = - fmap T.pack $ (:) <$> satisfy isAlpha <*> many (satisfy constituent) - where - constituent c = isAlphaNum c || c == '_' - lexeme :: Parser () -> Parser a -> Parser a lexeme sep p = p <* sep @@ -110,8 +120,9 @@ choice [ inParens sep (mkTuple <$> (parseExp sep `sepBy` pComma)), inBraces sep (Record <$> (pField `sepBy` pComma)), - Call <$> lexeme sep parseEntryName <*> many (parseExp sep), - Const <$> lexeme sep V.parsePrimValue + Call <$> lexeme sep parseFunc <*> many (parseExp sep), + Const <$> lexeme sep V.parsePrimValue, + StringLit . T.pack <$> lexeme sep ("\"" *> manyTill charLiteral "\"") ] where pField = (,) <$> parseEntryName <*> (pEquals *> parseExp sep) @@ -120,6 +131,17 @@ mkTuple [v] = v mkTuple vs = Tuple vs + parseFunc = + choice + [ FuncBuiltin <$> ("$" *> parseEntryName), + FuncFut <$> parseEntryName + ] + + parseEntryName = + fmap T.pack $ (:) <$> satisfy isAlpha <*> many (satisfy constituent) + where + constituent c = isAlphaNum c || c == '_' + prettyFailure :: CmdFailure -> T.Text prettyFailure (CmdFailure bef aft) = T.unlines $ bef ++ aft @@ -144,20 +166,24 @@ Just [val] -> pure $ Right val _ -> pure $ Left "Invalid data file produced by Futhark server." -writeVar :: (MonadError T.Text m, MonadIO m) => Server -> VarName -> PrimValue -> m () +writeVar :: (MonadError T.Text m, MonadIO m) => Server -> VarName -> V.Value -> m () writeVar server v val = cmdMaybe . liftIO . withSystemTempFile "futhark-server-write" $ \tmpf tmpf_h -> do - T.hPutStr tmpf_h $ prettyText val + LBS.hPutStr tmpf_h $ Bin.encode val hClose tmpf_h - let t = prettyText $ primValueType val - cmdRestore server tmpf [(v, t)] + -- We are not using prettyprinting for the type, because we don't + -- want the sizes of the dimensions. + let V.ValueType dims t = V.valueType val + t' = mconcat (map (const "[]") dims) <> prettyText t + cmdRestore server tmpf [(v, t')] -- | A ScriptValue is either a base value or a partially applied -- function. We don't have real first-class functions in -- FutharkScript, but we sort of have closures. data ScriptValue v = SValue TypeName v - | -- | Ins, then outs. + | -- | Ins, then outs. Yes, this is the opposite of more or less + -- everywhere else. SFun EntryName [TypeName] [TypeName] [ScriptValue v] instance Functor ScriptValue where @@ -176,7 +202,7 @@ = STValue TypeName | -- | Ins, then outs. STFun [TypeName] [TypeName] - deriving (Show) + deriving (Eq, Show) instance Pretty ScriptValueType where ppr (STValue t) = ppr t @@ -187,6 +213,13 @@ [out] -> strictText out _ -> parens $ commasep $ map strictText outs +data ValOrVar = VVal V.Value | VVar VarName + deriving (Show) + +-- | The intermediate values used during expression evaluation - in +-- particular, these may not be on the server. +type InterValue = V.Compound (ScriptValue ValOrVar) + -- | The value that is produced by expression evaluation. This -- representation keeps all values on the server. type ExpValue = V.Compound (ScriptValue VarName) @@ -207,15 +240,18 @@ valueToExp (V.ValueAtom (SValue t v)) = ServerVar t v valueToExp (V.ValueAtom (SFun fname _ _ closure)) = - Call fname $ map (valueToExp . V.ValueAtom) closure + Call (FuncFut fname) $ map (valueToExp . V.ValueAtom) closure valueToExp (V.ValueRecord fs) = Record $ M.toList $ M.map valueToExp fs valueToExp (V.ValueTuple fs) = Tuple $ map valueToExp fs +-- | How to evaluate a builtin function. +type EvalBuiltin m = T.Text -> [V.CompoundValue] -> m V.CompoundValue + -- | Evaluate a FutharkScript expression relative to some running server. -evalExp :: (MonadError T.Text m, MonadIO m) => ScriptServer -> Exp -> m ExpValue -evalExp (ScriptServer server counter) top_level_e = do +evalExp :: forall m. (MonadError T.Text m, MonadIO m) => EvalBuiltin m -> ScriptServer -> Exp -> m ExpValue +evalExp builtin (ScriptServer server counter) top_level_e = do vars <- liftIO $ newIORef [] let newVar base = liftIO $ do x <- readIORef counter @@ -224,38 +260,97 @@ modifyIORef vars (v :) pure v - evalExpToVar e = do - vs <- evalExpToVars e - case vs of - V.ValueAtom (SValue _ v) -> pure v - V.ValueAtom SFun {} -> - throwError $ "Expression " <> prettyText e <> " not fully applied." - _ -> - throwError $ "Expression " <> prettyText e <> " produced more than one value." - evalExpToVars (ServerVar t v) = - pure $ V.ValueAtom $ SValue t v - evalExpToVars (Call name es) = do - ins <- mapM evalExpToVar es + toVal :: ValOrVar -> m V.Value + toVal (VVal v) = pure v + toVal (VVar v) = readVar server v + + toVar :: ValOrVar -> m VarName + toVar (VVar v) = pure v + toVar (VVal val) = do + v <- newVar "const" + writeVar server v val + pure v + + scriptValueToValOrVar (SFun f _ _ _) = + throwError $ "Function " <> f <> " not fully applied." + scriptValueToValOrVar (SValue _ v) = + pure v + + scriptValueToVal :: ScriptValue ValOrVar -> m V.Value + scriptValueToVal = toVal <=< scriptValueToValOrVar + + scriptValueToVar :: ScriptValue ValOrVar -> m VarName + scriptValueToVar = toVar <=< scriptValueToValOrVar + + interValToVal :: InterValue -> m V.CompoundValue + interValToVal = traverse scriptValueToVal + + interValToVar :: InterValue -> m VarName + interValToVar (V.ValueAtom v) = scriptValueToVar v + interValToVar _ = throwError "Unexpected tuple or record value." + + valToInterVal :: V.CompoundValue -> InterValue + valToInterVal = fmap $ \v -> + SValue (prettyText (V.valueType v)) $ VVal v + + interValToExpVal :: InterValue -> m ExpValue + interValToExpVal = traverse (traverse toVar) + + simpleType (V.ValueAtom (STValue _)) = True + simpleType _ = False + + evalExp' :: Exp -> m InterValue + evalExp' (ServerVar t v) = + pure $ V.ValueAtom $ SValue t $ VVar v + evalExp' (Call (FuncBuiltin name) es) = do + v <- builtin name =<< mapM (interValToVal <=< evalExp') es + pure $ valToInterVal v + evalExp' (Call (FuncFut name) es) = do in_types <- cmdEither $ cmdInputs server name out_types <- cmdEither $ cmdOutputs server name + + es' <- mapM evalExp' es + let es_types = map (fmap scriptValueType) es' + + unless (all simpleType es_types) $ + throwError $ + "Literate Futhark does not support passing script-constructed records, tuples, or functions to entry points.\n" + <> "Create a Futhark wrapper function." + + -- Careful to not require saturated application. + unless (and $ zipWith (==) es_types (map (V.ValueAtom . STValue) in_types)) $ + throwError $ + "Function \"" <> name <> "\" expects arguments of types:\n" + <> prettyText (V.ValueTuple $ map V.ValueAtom in_types) + <> "\nBut called with arguments of types:\n" + <> prettyText (V.ValueTuple $ map V.ValueAtom es_types) + + ins <- mapM (interValToVar <=< evalExp') es + if length in_types == length ins then do outs <- replicateM (length out_types) $ newVar "out" void $ cmdEither $ cmdCall server name outs ins - pure $ V.mkCompound $ zipWith SValue out_types outs + pure $ V.mkCompound $ zipWith SValue out_types $ map VVar outs else pure . V.ValueAtom . SFun name in_types out_types $ - zipWith SValue in_types ins - evalExpToVars (Const val) = do - v <- newVar "const" - writeVar server v val - pure $ V.ValueAtom $ SValue (prettyText (primValueType val)) v - evalExpToVars (Tuple es) = - V.ValueTuple <$> mapM evalExpToVars es - evalExpToVars e@(Record m) = do + zipWith SValue in_types $ map VVar ins + evalExp' (StringLit s) = + case V.putValue s of + Just s' -> + pure $ V.ValueAtom $ SValue (prettyText (V.valueType s')) $ VVal s' + Nothing -> error $ "Unable to write value " ++ pretty s + evalExp' (Const val) = + case V.putValue val of + Just val' -> + pure $ V.ValueAtom $ SValue (prettyText (primValueType val)) $ VVal val' + Nothing -> error $ "Unable to write value " ++ pretty val + evalExp' (Tuple es) = + V.ValueTuple <$> mapM evalExp' es + evalExp' e@(Record m) = do when (length (nubOrd (map fst m)) /= length (map fst m)) $ throwError $ "Record " <> prettyText e <> " has duplicate fields." - V.ValueRecord <$> traverse evalExpToVars (M.fromList m) + V.ValueRecord <$> traverse evalExp' (M.fromList m) let freeNonresultVars v = do let v_vars = serverVarsInValue v @@ -270,12 +365,12 @@ -- Call. void $ liftIO $ cmdFree server =<< readIORef vars throwError e - (freeNonresultVars =<< evalExpToVars top_level_e) `catchError` freeVarsOnError + (freeNonresultVars =<< interValToExpVal =<< evalExp' top_level_e) `catchError` freeVarsOnError -- | Read actual values from the server. Fails for values that have -- no well-defined external representation. getExpValue :: - (MonadError T.Text m, MonadIO m) => ScriptServer -> ExpValue -> m (V.Compound V.Value) + (MonadError T.Text m, MonadIO m) => ScriptServer -> ExpValue -> m V.CompoundValue getExpValue (ScriptServer server _) e = traverse toGround =<< traverse (traverse (readVar server)) e where @@ -285,18 +380,20 @@ -- | Like 'evalExp', but requires all values to be non-functional. evalExpToGround :: - (MonadError T.Text m, MonadIO m) => ScriptServer -> Exp -> m (V.Compound V.Value) -evalExpToGround server e = getExpValue server =<< evalExp server e + (MonadError T.Text m, MonadIO m) => EvalBuiltin m -> ScriptServer -> Exp -> m V.CompoundValue +evalExpToGround builtin server e = getExpValue server =<< evalExp builtin server e -- | The set of Futhark variables that are referenced by the -- expression - these will have to be entry points in the Futhark -- program. varsInExp :: Exp -> S.Set EntryName varsInExp ServerVar {} = mempty -varsInExp (Call v es) = S.insert v $ foldMap varsInExp es +varsInExp (Call (FuncFut v) es) = S.insert v $ foldMap varsInExp es +varsInExp (Call (FuncBuiltin _) es) = foldMap varsInExp es varsInExp (Tuple es) = foldMap varsInExp es varsInExp (Record fs) = foldMap (foldMap varsInExp) fs varsInExp Const {} = mempty +varsInExp StringLit {} = mempty -- | Release all the server-side variables in the value. Yes, -- FutharkScript has manual memory management... diff --git a/src/Futhark/Server.hs b/src/Futhark/Server.hs --- a/src/Futhark/Server.hs +++ b/src/Futhark/Server.hs @@ -58,15 +58,28 @@ case code of Just (ExitFailure e) -> error $ "Cannot start " ++ prog ++ ": error " ++ show e - _ -> - pure $ - Server - { serverStdin = stdin, - serverStdout = stdout, - serverProc = phandle, - serverDebug = isEnvVarAtLeast "FUTHARK_COMPILER_DEBUGGING" 1, - serverErrLog = err_log_f - } + _ -> do + let server = + Server + { serverStdin = stdin, + serverStdout = stdout, + serverProc = phandle, + serverDebug = isEnvVarAtLeast "FUTHARK_COMPILER_DEBUGGING" 1, + serverErrLog = err_log_f + } + void (responseLines server) `catch` onStartupError server + pure server + where + onStartupError :: Server -> IOError -> IO a + onStartupError s _ = do + code <- P.waitForProcess $ serverProc s + stderr_s <- readFile $ serverErrLog s + removeFile $ serverErrLog s + error $ + "Command failed with " ++ show code ++ ":\n" + ++ unwords (prog : options) + ++ "\nStderr:\n" + ++ stderr_s stopServer :: Server -> IO () stopServer s = do diff --git a/src/Futhark/Test/Values.hs b/src/Futhark/Test/Values.hs --- a/src/Futhark/Test/Values.hs +++ b/src/Futhark/Test/Values.hs @@ -49,6 +49,7 @@ import Data.Int (Int16, Int32, Int64, Int8) import qualified Data.Map as M import qualified Data.Text as T +import qualified Data.Text.Encoding as T import Data.Traversable import Data.Vector.Generic (freeze) import qualified Data.Vector.Storable as SVec @@ -726,6 +727,9 @@ class GetValue t where getValue :: Value -> Maybe t +instance GetValue t => GetValue [t] where + getValue = mapM getValue . valueElems + instance GetValue Bool where getValue (BoolValue shape vs) | [] <- SVec.toList shape = @@ -756,12 +760,39 @@ Just $ vs SVec.! 0 getValue _ = Nothing +instance GetValue Word8 where + getValue (Word8Value shape vs) + | [] <- SVec.toList shape = + Just $ vs SVec.! 0 + getValue _ = Nothing + +instance GetValue Word16 where + getValue (Word16Value shape vs) + | [] <- SVec.toList shape = + Just $ vs SVec.! 0 + getValue _ = Nothing + +instance GetValue Word32 where + getValue (Word32Value shape vs) + | [] <- SVec.toList shape = + Just $ vs SVec.! 0 + getValue _ = Nothing + +instance GetValue Word64 where + getValue (Word64Value shape vs) + | [] <- SVec.toList shape = + Just $ vs SVec.! 0 + getValue _ = Nothing + -- | A class for Haskell values that can be converted to 'Value'. -- This is a convenience facility - don't expect it to be fast. class PutValue t where -- | This may fail for cases such as irregular arrays. putValue :: t -> Maybe Value +instance PutValue Word8 where + putValue = Just . Word8Value mempty . SVec.singleton + instance PutValue F.PrimValue where putValue (F.SignedValue (F.Int8Value x)) = Just $ Int8Value mempty $ SVec.singleton x @@ -815,3 +846,12 @@ getVec (Float32Value _ vec) = SVec.unsafeCast vec getVec (Float64Value _ vec) = SVec.unsafeCast vec getVec (BoolValue _ vec) = SVec.unsafeCast vec + +instance PutValue T.Text where + putValue = putValue . T.encodeUtf8 + +instance PutValue BS.ByteString where + putValue bs = + Just $ Word8Value size $ byteStringToVector bs + where + size = SVec.fromList [fromIntegral (BS.length bs)] 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 @@ -20,6 +20,7 @@ import Control.Monad.Except import Control.Monad.State +import Data.Function ((&)) import Data.List (zip4) import qualified Data.Map.Strict as M import qualified Futhark.IR as AST @@ -196,13 +197,12 @@ (++ patternNames pat) <$> replicateM (length scanacc_params) (newVName "discard") letBindNames names $ DoLoop [] merge loopform loop_body -transformSOAC pat (Stream w stream_form lam arrs) = do +transformSOAC pat (Stream w _ lam nes arrs) = do -- Create a loop that repeatedly applies the lambda body to a -- chunksize of 1. Hopefully this will lead to this outer loop -- being the only one, as all the innermost one can be simplified -- array (as they will have one iteration each). - let nes = getStreamAccums stream_form - (chunk_size_param, fold_params, chunk_params) = + let (chunk_size_param, fold_params, chunk_params) = partitionChunkedFoldParameters (length nes) $ lambdaParams lam mapout_merge <- forM (drop (length nes) $ lambdaReturnType lam) $ \t -> @@ -255,11 +255,11 @@ transformSOAC pat (Scatter len lam ivs as) = do iter <- newVName "write_iter" - let (_as_ws, as_ns, as_vs) = unzip3 as + let (as_ws, as_ns, as_vs) = unzip3 as ts <- mapM lookupType as_vs asOuts <- mapM (newIdent "write_out") ts - let ivsLen = length (lambdaReturnType lam) `div` 2 + let ivsLen = zipWith (*) as_ns $ map length as_ws -- Scatter is in-place, so we use the input array as the output array. let merge = loopMerge asOuts $ map Var as_vs @@ -274,12 +274,15 @@ letSubExp "write_iv" $ BasicOp $ Index iv $ fullSlice iv_t [DimFix $ Var iter] ivs'' <- bindLambda lam (map (BasicOp . SubExp) ivs') - let indexes = chunks as_ns $ take ivsLen ivs'' - values = chunks as_ns $ drop ivsLen ivs'' + let indexes = + take (sum ivsLen) ivs'' + & chunks (concat $ zipWith (\ws n -> replicate n (length ws)) as_ws as_ns) + & chunks as_ns + values = chunks as_ns $ drop (sum ivsLen) ivs'' ress <- forM (zip3 indexes values (map identName asOuts)) $ \(indexes', values', arr) -> do let saveInArray arr' (indexCur, valueCur) = - letExp "write_out" =<< eWriteArray arr' [eSubExp indexCur] (eSubExp valueCur) + letExp "write_out" =<< eWriteArray arr' (map eSubExp indexCur) (eSubExp valueCur) foldM saveInArray arr $ zip indexes' values' return $ resultBody (map Var ress) 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 @@ -30,7 +30,7 @@ import Control.Monad.State import Control.Monad.Trans.Maybe import Data.Array -import Data.Bifunctor (first) +import Data.Bifunctor (first, second) import Data.List ( find, foldl', @@ -1004,11 +1004,11 @@ eval env (Lambda ps body _ (Info (_, rt)) _) = evalFunction env [] ps body rt eval env (OpSection qv (Info t) _) = evalTermVar env qv $ toStruct t -eval env (OpSectionLeft qv _ e (Info (_, argext), _) (Info t, Info retext) loc) = do +eval env (OpSectionLeft qv _ e (Info (_, _, argext), _) (Info t, Info retext) loc) = do v <- evalArg env e argext f <- evalTermVar env qv (toStruct t) returned env t retext =<< apply loc env f v -eval env (OpSectionRight qv _ e (Info _, Info (_, argext)) (Info t) loc) = do +eval env (OpSectionRight qv _ e (Info _, Info (_, _, argext)) (Info t) loc) = do y <- evalArg env e argext return $ ValueFun $ \x -> do @@ -1110,19 +1110,27 @@ env' <- patternMatch env p v lift $ eval env' cExp +-- We hackily do multiple substitutions in modules, because otherwise +-- we would lose in cases where the parameter substitutions are [a->x, +-- b->x] when we reverse. (See issue #1250.) +reverseSubstitutions :: M.Map VName VName -> M.Map VName [VName] +reverseSubstitutions = + M.fromListWith (<>) . map (second pure . uncurry (flip (,))) . M.toList + substituteInModule :: M.Map VName VName -> Module -> Module substituteInModule substs = onModule where rev_substs = reverseSubstitutions substs - replace v = fromMaybe v $ M.lookup v rev_substs - replaceQ v = maybe v qualName $ M.lookup (qualLeaf v) rev_substs + replace v = fromMaybe [v] $ M.lookup v rev_substs + replaceQ v = maybe v qualName $ maybeHead =<< M.lookup (qualLeaf v) rev_substs replaceM f m = M.fromList $ do (k, v) <- M.toList m - return (replace k, f v) + k' <- replace k + return (k', f v) onModule (Module (Env terms types _)) = Module $ Env (replaceM onTerm terms) (replaceM onType types) mempty onModule (ModuleFun f) = - ModuleFun $ \m -> onModule <$> f (substituteInModule rev_substs m) + ModuleFun $ \m -> onModule <$> f (substituteInModule (M.mapMaybe maybeHead rev_substs) m) onTerm (TermValue t v) = TermValue t v onTerm (TermPoly t v) = TermPoly t v onTerm (TermModule m) = TermModule $ onModule m @@ -1131,9 +1139,6 @@ onDim (ConstDim x) = ConstDim x onDim AnyDim = AnyDim -reverseSubstitutions :: M.Map VName VName -> M.Map VName VName -reverseSubstitutions = M.fromList . map (uncurry $ flip (,)) . M.toList - evalModuleVar :: Env -> QualName VName -> EvalM Module evalModuleVar env qv = case lookupVar qv env of @@ -1555,6 +1560,36 @@ if i >= 0 && i < arrayLength arr' then arr' // [(i, v)] else arr' + def "scatter_2d" = Just $ + fun3t $ \arr is vs -> + case arr of + ValueArray _ _ -> + return $ + foldl' update arr $ + zip (map fromTuple $ snd $ fromArray is) (snd $ fromArray vs) + _ -> + error $ "scatter_2d expects array, but got: " ++ pretty arr + where + update :: Value -> (Maybe [Value], Value) -> Value + update arr (Just idxs@[_, _], v) = + fromMaybe arr $ updateArray (map (IndexingFix . asInt64) idxs) arr v + update _ _ = + error "scatter_2d expects 2-dimensional indices" + def "scatter_3d" = Just $ + fun3t $ \arr is vs -> + case arr of + ValueArray _ _ -> + return $ + foldl' update arr $ + zip (map fromTuple $ snd $ fromArray is) (snd $ fromArray vs) + _ -> + error $ "scatter_3d expects array, but got: " ++ pretty arr + where + update :: Value -> (Maybe [Value], Value) -> Value + update arr (Just idxs@[_, _, _], v) = + fromMaybe arr $ updateArray (map (IndexingFix . asInt64) idxs) arr v + update _ _ = + error "scatter_3d expects 3-dimensional indices" def "hist" = Just $ fun6t $ \_ arr fun _ is vs -> case arr of diff --git a/src/Language/Futhark/Pretty.hs b/src/Language/Futhark/Pretty.hs --- a/src/Language/Futhark/Pretty.hs +++ b/src/Language/Futhark/Pretty.hs @@ -331,9 +331,9 @@ pprPrec _ (OpSection binop _ _) = parens $ ppr binop pprPrec _ (OpSectionLeft binop _ x _ _ _) = - parens $ ppr x <+> ppr binop + parens $ ppr x <+> ppBinOp binop pprPrec _ (OpSectionRight binop _ x _ _ _) = - parens $ ppr binop <+> ppr x + parens $ ppBinOp binop <+> ppr x pprPrec _ (ProjectSection fields _ _) = parens $ mconcat $ map p fields where @@ -495,6 +495,15 @@ sig' = case sig of Nothing -> mempty Just (s, _) -> colon <+> ppr s <> text " " + +ppBinOp :: IsName v => QualName v -> Doc +ppBinOp bop = + case leading of + Backtick -> text "`" <> ppr bop <> text "`" + _ -> ppr bop + where + leading = + leadingOperator $ nameFromString $ pretty $ pprName $ qualLeaf bop prettyBinOp :: (Eq vn, IsName vn, Annot f) => 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 @@ -624,10 +624,10 @@ named (Unnamed, _) = Nothing typeOf (OpSection _ (Info t) _) = t -typeOf (OpSectionLeft _ _ _ (_, Info pt2) (Info ret, _) _) = - foldFunType [fromStruct pt2] ret -typeOf (OpSectionRight _ _ _ (Info pt1, _) (Info ret) _) = - foldFunType [fromStruct pt1] ret +typeOf (OpSectionLeft _ _ _ (_, Info (pn, pt2)) (Info ret, _) _) = + Scalar $ Arrow mempty pn (fromStruct pt2) ret +typeOf (OpSectionRight _ _ _ (Info (pn, pt1), _) (Info ret) _) = + Scalar $ Arrow mempty pn (fromStruct pt1) ret typeOf (ProjectSection _ (Info t) _) = t typeOf (IndexSection _ (Info t) _) = t typeOf (Constr _ _ (Info t) _) = t @@ -879,6 +879,24 @@ ] $ Array () Unique t_a (rank 1) ), + ( "scatter_2d", + IntrinsicPolyFun + [tp_a] + [ uarr_2d_a, + Array () Nonunique (tupInt64 2) (rank 1), + Array () Nonunique t_a (rank 1) + ] + uarr_2d_a + ), + ( "scatter_3d", + IntrinsicPolyFun + [tp_a] + [ uarr_3d_a, + Array () Nonunique (tupInt64 3) (rank 1), + Array () Nonunique t_a (rank 1) + ] + uarr_3d_a + ), ("zip", IntrinsicPolyFun [tp_a, tp_b] [arr_a, arr_b] arr_a_b), ("unzip", IntrinsicPolyFun [tp_a, tp_b] [arr_a_b] t_arr_a_arr_b), ( "hist", @@ -959,6 +977,8 @@ t_a = TypeVar () Nonunique (typeName tv_a) [] arr_a = Array () Nonunique t_a (rank 1) arr_2d_a = Array () Nonunique t_a (rank 2) + uarr_2d_a = Array () Unique t_a (rank 2) + uarr_3d_a = Array () Unique t_a (rank 3) uarr_a = Array () Unique t_a (rank 1) tp_a = TypeParamType Unlifted tv_a mempty @@ -1053,6 +1073,12 @@ intrinsicBinOp Greater = ordering intrinsicBinOp Geq = ordering intrinsicBinOp _ = Nothing + + tupInt64 n = + Record $ + M.fromList $ + zip tupleFieldNames $ + replicate n $ Scalar $ Prim $ Signed Int64 -- | The largest tag used by an intrinsic - this can be used to -- determine whether a 'VName' refers to an intrinsic or a user-defined name. 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 @@ -128,6 +128,8 @@ Show (f Int), Show (f StructType), Show (f (StructType, Maybe VName)), + Show (f (PName, StructType)), + Show (f (PName, StructType, Maybe VName)), Show (f (Aliasing, StructType)), Show (f (M.Map VName VName)), Show (f Uniqueness) @@ -753,7 +755,7 @@ (QualName vn) (f PatternType) (ExpBase f vn) - (f (StructType, Maybe VName), f StructType) + (f (PName, StructType, Maybe VName), f (PName, StructType)) (f PatternType, f [VName]) SrcLoc | -- | @+2@; first type is operand, second is result. @@ -761,7 +763,7 @@ (QualName vn) (f PatternType) (ExpBase f vn) - (f StructType, f (StructType, Maybe VName)) + (f (PName, StructType), f (PName, StructType, Maybe VName)) (f PatternType) SrcLoc | -- | Field projection as a section: @(.x.y.z)@. diff --git a/src/Language/Futhark/Traversals.hs b/src/Language/Futhark/Traversals.hs --- a/src/Language/Futhark/Traversals.hs +++ b/src/Language/Futhark/Traversals.hs @@ -1,4 +1,5 @@ {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TupleSections #-} -- | -- @@ -169,23 +170,23 @@ OpSection <$> mapOnQualName tv name <*> traverse (mapOnPatternType tv) t <*> pure loc - astMap tv (OpSectionLeft name t arg (Info (t1a, argext), t1b) (t2, retext) loc) = + astMap tv (OpSectionLeft name t arg (Info (pa, t1a, argext), Info (pb, t1b)) (t2, retext) loc) = OpSectionLeft <$> mapOnQualName tv name <*> traverse (mapOnPatternType tv) t <*> mapOnExp tv arg <*> ( (,) - <$> (Info <$> ((,) <$> mapOnStructType tv t1a <*> pure argext)) - <*> traverse (mapOnStructType tv) t1b + <$> (Info <$> ((pa,,) <$> mapOnStructType tv t1a <*> pure argext)) + <*> (Info <$> ((pb,) <$> mapOnStructType tv t1b)) ) <*> ((,) <$> traverse (mapOnPatternType tv) t2 <*> pure retext) <*> pure loc - astMap tv (OpSectionRight name t arg (t1a, Info (t1b, argext)) t2 loc) = + astMap tv (OpSectionRight name t arg (Info (pa, t1a), Info (pb, t1b, argext)) t2 loc) = OpSectionRight <$> mapOnQualName tv name <*> traverse (mapOnPatternType tv) t <*> mapOnExp tv arg <*> ( (,) - <$> traverse (mapOnStructType tv) t1a - <*> (Info <$> ((,) <$> mapOnStructType tv t1b <*> pure argext)) + <$> (Info <$> ((pa,) <$> mapOnStructType tv t1a)) + <*> (Info <$> ((pb,,) <$> mapOnStructType tv t1b <*> pure argext)) ) <*> traverse (mapOnPatternType tv) t2 <*> pure loc diff --git a/src/Language/Futhark/TypeChecker/Terms.hs b/src/Language/Futhark/TypeChecker/Terms.hs --- a/src/Language/Futhark/TypeChecker/Terms.hs +++ b/src/Language/Futhark/TypeChecker/Terms.hs @@ -1668,14 +1668,14 @@ (op', ftype) <- lookupVar loc op e_arg <- checkArg e (t1, rt, argext, retext) <- checkApply loc (Just op', 0) ftype e_arg - case rt of - Scalar (Arrow _ _ t2 rettype) -> + case (ftype, rt) of + (Scalar (Arrow _ m1 _ _), Scalar (Arrow _ m2 t2 rettype)) -> return $ OpSectionLeft op' (Info ftype) (argExp e_arg) - (Info (toStruct t1, argext), Info $ toStruct t2) + (Info (m1, toStruct t1, argext), Info (m2, toStruct t2)) (Info rettype, Info retext) loc _ -> @@ -1697,7 +1697,7 @@ op' (Info ftype) (argExp e_arg) - (Info $ toStruct t1, Info (toStruct t2', argext)) + (Info (m1, toStruct t1), Info (m2, toStruct t2', argext)) (Info $ addAliases ret (<> aliases ret')) loc _ ->